├── README.md ├── batch_rename.py ├── combine_dataset.py ├── img_cvt_format.py ├── interactive_image_segmentation.py ├── resouce └── images │ ├── 000009.png │ ├── 000233.png │ ├── 000457.png │ ├── 000681.png │ ├── 000905.png │ ├── 001129.png │ └── 001353.png ├── test_grabcut.py └── video2img.py /README.md: -------------------------------------------------------------------------------- 1 | # InteractiveImageSegmentation 2 | An interactive image segmentation tool for pixel-wise labeling image dataset in segmentation task, 3 | which use GrabCut("[“GrabCut”: interactive foreground extraction using iterated graph cuts](https://dl.acm.org/citation.cfm?id=1015720)") and implemented in OpenCV 3 and Python. 4 | An OpenCV tutorial can be found [here](https://docs.opencv.org/3.0-beta/doc/py_tutorials/py_imgproc/py_grabcut/py_grabcut.html?highlight=grabcut). 5 | 6 | ## Quick Start 7 | Go to this project dir and run 'python interactive_image_segmentation.py resouce/images/ resouce/labels'. You will see a UI window and you can start labelling. 8 | If you want to label your own images, collect them into a filefolder, modify the source_dir and save_dir. 9 | 10 | ## How to use 11 | - CTRL+left mouse button: label certain background pixels 12 | - SHIFT+left mouse button: label certain foreground pixels 13 | - CTRL+right mouse button: label possible background pixels 14 | - SHIFT+right mouse button: label possible foreground pixels 15 | - 'a'/SPACE: run sengementation again 16 | - 's'/ENTER: save label and skip to next image 17 | - 'p': prev image 18 | - 'n': next image 19 | - 'q'/ESC: exit 20 | 21 | There are some useful py scripts you can use to handle your own data. Usage: 22 | - 'interactive_image_segmentation.py image_dir save_dir' runs interactive image segmentation 23 | - 'batch_rename.py dir1 [dir2] ...' renames all images in folder with sequencial numbers. 24 | - 'img_cvt_format.py image_dir [suffix]' converts all images in folder to the same image format, such as '.png' '.jpg'... 25 | - 'video2img.py video_file save_dir [stride] [resize] [suffix]' converts a video file to a list of images. 26 | 27 | ## License 28 | This project is released under a MIT license. 29 | -------------------------------------------------------------------------------- /batch_rename.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | 4 | def batch_rename(filedir): 5 | idx = 0 6 | for f in sorted([x for x in os.listdir(filedir)]): 7 | os.rename(os.path.join(filedir,f),os.path.join(filedir,'%06d.'%idx+f.rsplit('.')[1])) 8 | idx+=1 9 | 10 | if __name__ == '__main__': 11 | if len(sys.argv)<2 or '-h' in sys.argv[1]: 12 | print('Usage: batch_rename [datadir] [datadir2] ...') 13 | else: 14 | for i in range(1,len(sys.argv)): 15 | batch_rename(sys.argv[i]) -------------------------------------------------------------------------------- /combine_dataset.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import cv2 3 | import numpy as np 4 | import sys 5 | import os 6 | 7 | img_size = (480,270) 8 | save_dir = '/home/symao/data/skyseg' 9 | # list of (data_dir, downsample, augment_rate) 10 | data_dirs = [('/home/symao/data/NVIDIA-Aerial-Drone-Dataset/FPV/AT/GOPR2188',1, 0.3), 11 | ('/home/symao/data/NVIDIA-Aerial-Drone-Dataset/FPV/AT/GOPR2189',1, 3), 12 | ('/home/symao/data/NVIDIA-Aerial-Drone-Dataset/FPV/SFWA/360p',8, 0), 13 | ('/home/symao/data/NVIDIA-Aerial-Drone-Dataset/internet_pictures',1, 4)] 14 | save_idx = 0 15 | 16 | def img_diversity_illu(img): 17 | img = np.copy(img) 18 | illu_M_foo = [] 19 | illu_M_foo.append(lambda x:np.ones(x.shape,dtype='uint8')*np.random.randint(30,50)) 20 | illu_M_foo.append(lambda x:np.ones(x.shape,dtype='uint8')*np.random.randint(60,80)) 21 | illu_M_foo.append(lambda x:np.ones(x.shape,dtype='uint8')*((np.ones(img.shape[1:],dtype='uint8').T*(np.linspace(np.random.randint(0,10), np.random.randint(30,60),img.shape[1])).astype('uint8')).T)) 22 | illu_M_foo.append(lambda x:np.ones(x.shape,dtype='uint8')*((np.ones(img.shape[1:],dtype='uint8').T*(np.linspace(np.random.randint(50,90), np.random.randint(0,30),img.shape[1])).astype('uint8')).T)) 23 | 24 | illu_foo = [] 25 | illu_foo.append(lambda x,M:cv2.add(x,M)) 26 | illu_foo.append(lambda x,M:cv2.subtract(x,M)) 27 | return np.random.choice(illu_foo)(img,np.random.choice(illu_M_foo)(img)) 28 | 29 | def save_img(img_dir, label_dir, save_dir, downsample = 1, augment = 0): 30 | global save_idx 31 | saveimgdir = os.path.join(save_dir,'images') 32 | savelbldir = os.path.join(save_dir,'labels') 33 | if not os.path.exists(saveimgdir): 34 | os.makedirs(saveimgdir) 35 | if not os.path.exists(savelbldir): 36 | os.makedirs(savelbldir) 37 | 38 | fimglist = sorted([x for x in os.listdir(img_dir) if '.png' in x or '.jpg' in x]) 39 | for i in range(0,len(fimglist),downsample): 40 | fimg = fimglist[i] 41 | if os.path.exists(os.path.join(label_dir,fimg)): 42 | img = cv2.resize(cv2.imread(os.path.join(img_dir,fimg)),img_size) 43 | lbl = cv2.resize(cv2.imread(os.path.join(label_dir,fimg)),img_size) 44 | cv2.imwrite(os.path.join(saveimgdir,'%06d.png'%save_idx),img) 45 | cv2.imwrite(os.path.join(savelbldir,'%06d.png'%save_idx),lbl) 46 | save_idx+=1 47 | print('saving...%d/%d'%(i+1,len(fimglist)), end='\r') 48 | print('') 49 | if augment>0: 50 | print('augmenting...') 51 | for i in range(0,len(fimglist),downsample): 52 | fimg = fimglist[i] 53 | if os.path.exists(os.path.join(label_dir,fimg)): 54 | img = cv2.resize(cv2.imread(os.path.join(img_dir,fimg)),img_size) 55 | lbl = cv2.resize(cv2.imread(os.path.join(label_dir,fimg)),img_size) 56 | temp_rate = augment 57 | while temp_rate>0: 58 | if np.random.rand()0 and self.img.size>0: 56 | if flags & cv2.EVENT_FLAG_CTRLKEY: 57 | cv2.circle(self.img, (x,y), self.radius, (COLOR_BG if self.left_mouse_down else tuple([k/3 for k in COLOR_BG])), -1) 58 | cv2.circle(self.mask, (x,y), self.radius, (cv2.GC_BGD if self.left_mouse_down else cv2.GC_PR_BGD), -1) 59 | elif flags & cv2.EVENT_FLAG_SHIFTKEY: 60 | cv2.circle(self.img, (x,y), self.radius, (COLOR_FG if self.left_mouse_down else tuple([k/3 for k in COLOR_FG])), -1) 61 | cv2.circle(self.mask, (x,y), self.radius, (cv2.GC_FGD if self.left_mouse_down else cv2.GC_PR_FGD), -1) 62 | if event == cv2.EVENT_MOUSEWHEEL: 63 | if flags<0: 64 | diff_k = int(np.clip(self.radius*0.4,1,5)) 65 | self.radius+=diff_k 66 | elif flags>0: 67 | diff_k = int(np.clip(self.radius*0.4,1,5)) 68 | self.radius-=diff_k 69 | self.radius = np.clip(self.radius, 1, self.max_radius) 70 | cv2.setTrackbarPos('brush size', self.winname, self.radius) 71 | 72 | def __init_mask(self, mask): 73 | mask[:] = cv2.GC_PR_FGD 74 | mask[:10,:] = cv2.GC_PR_BGD 75 | 76 | def process(self, img): 77 | self.img = np.copy(img) 78 | if self.use_prev_mask==False or self.mask.shape[:2]!=self.img.shape[:2]: 79 | self.mask = np.zeros(img.shape[:2],'uint8') 80 | self.__init_mask(self.mask) 81 | self.bgdModel = np.zeros((1,65),np.float64) 82 | self.fgdModel = np.zeros((1,65),np.float64) 83 | cv2.grabCut(img, self.mask, None, self.bgdModel, self.fgdModel, 1, cv2.GC_INIT_WITH_MASK) 84 | 85 | while True: 86 | self.radius = cv2.getTrackbarPos('brush size',self.winname) 87 | color = mask2color(self.mask) 88 | alpha = 0.5 if self.draw_color==0 else (1 if self.draw_color==1 else 0) 89 | show_img = (self.img*alpha + color*(1-alpha)).astype('uint8') 90 | cv2.circle(show_img, self.cur_mouse, self.radius, (200,200,200), (2 if self.left_mouse_down else 1)) 91 | cv2.imshow(self.winname,show_img) 92 | cv2.imshow('color',color) 93 | key = cv2.waitKey(100) 94 | if key == ord('c'): 95 | self.img = np.copy(img) 96 | self.__init_mask(self.mask) 97 | elif key == ord('q') or key == 27 or key==ord('s') or key==ord('p') or key==ord('n') or key == 10: 98 | break 99 | elif key == ord('w'): 100 | self.draw_color = (self.draw_color+1)%3 101 | elif key == ord('a') or key == 32: 102 | cv2.putText(show_img, 'segmenting...', (10,30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,255,255),2) 103 | cv2.imshow(self.winname,show_img) 104 | cv2.waitKey(1) 105 | # mask enum 106 | # GC_BGD = 0, //背景 107 | # GC_FGD = 1, //前景 108 | # GC_PR_BGD = 2, //可能背景 109 | # GC_PR_FGD = 3 //可能前景 110 | hist, _ = np.histogram(self.mask,[0,1,2,3,4]) 111 | if hist[0]+hist[2]!=0 and hist[1]+hist[3]!=0: 112 | cv2.grabCut(img, self.mask, None, self.bgdModel, self.fgdModel, 1, cv2.GC_INIT_WITH_MASK) 113 | self.img = np.copy(img) 114 | return key 115 | 116 | if __name__ == '__main__': 117 | if(len(sys.argv)!=3): 118 | print('Usage: interactive_image_segmentation.py [img_dir] [save_dir]') 119 | exit() 120 | 121 | img_dir = sys.argv[1] 122 | save_dir = sys.argv[2] 123 | 124 | if not os.path.exists(save_dir): 125 | os.makedirs(save_dir) 126 | print('%s not exists, create it.'%save_dir) 127 | 128 | print("================= Interactive Image Segmentation =================") 129 | print("CTRL+left mouse button: select certain background pixels ") 130 | print("SHIFT+left mouse button: select certain foreground pixels ") 131 | print("CTRL+right mouse button: select possible background pixels ") 132 | print("SHIFT+right mouse button: select possible foreground pixels ") 133 | print("'a'/SPACE: run sengementation again") 134 | print("'p': prev image 'n': next image") 135 | print("'s'/ENTER: save label 'q'/ESC: exit") 136 | 137 | iis = InteractiveImageSegmentation() 138 | iis.use_prev_mask = True 139 | fimglist = sorted([x for x in os.listdir(img_dir) if '.png' in x or '.jpg' in x]) 140 | idx = 0 141 | while idx0: 156 | idx -= 1 157 | elif key == ord('n') or key == 32: 158 | idx += 1 159 | elif key == ord('q') or key == 27: 160 | break 161 | iis.mask[np.where(iis.mask==cv2.GC_BGD)]=cv2.GC_PR_BGD 162 | iis.mask[np.where(iis.mask==cv2.GC_FGD)]=cv2.GC_PR_FGD 163 | -------------------------------------------------------------------------------- /resouce/images/000009.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/symao/InteractiveImageSegmentation/b5e141fb6b8ea280a5e2f9b173bc37f25ff9e90d/resouce/images/000009.png -------------------------------------------------------------------------------- /resouce/images/000233.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/symao/InteractiveImageSegmentation/b5e141fb6b8ea280a5e2f9b173bc37f25ff9e90d/resouce/images/000233.png -------------------------------------------------------------------------------- /resouce/images/000457.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/symao/InteractiveImageSegmentation/b5e141fb6b8ea280a5e2f9b173bc37f25ff9e90d/resouce/images/000457.png -------------------------------------------------------------------------------- /resouce/images/000681.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/symao/InteractiveImageSegmentation/b5e141fb6b8ea280a5e2f9b173bc37f25ff9e90d/resouce/images/000681.png -------------------------------------------------------------------------------- /resouce/images/000905.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/symao/InteractiveImageSegmentation/b5e141fb6b8ea280a5e2f9b173bc37f25ff9e90d/resouce/images/000905.png -------------------------------------------------------------------------------- /resouce/images/001129.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/symao/InteractiveImageSegmentation/b5e141fb6b8ea280a5e2f9b173bc37f25ff9e90d/resouce/images/001129.png -------------------------------------------------------------------------------- /resouce/images/001353.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/symao/InteractiveImageSegmentation/b5e141fb6b8ea280a5e2f9b173bc37f25ff9e90d/resouce/images/001353.png -------------------------------------------------------------------------------- /test_grabcut.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import sys 4 | import os 5 | 6 | img = cv2.imread('/home/symao/Pictures/3773.png') 7 | mask = np.ones(img.shape[:2],np.uint8)*3 8 | bgdModel = np.zeros((1,65),np.float64) 9 | fgdModel = np.zeros((1,65),np.float64) 10 | mask[:3,:] = cv2.GC_PR_BGD 11 | cv2.grabCut(img,mask,None,bgdModel,fgdModel,1,cv2.GC_INIT_WITH_MASK) 12 | mask2 = np.where((mask==2)|(mask==0),0,1).astype('uint8') 13 | img = img*mask2[:,:,np.newaxis] 14 | 15 | cv2.imshow('img',img) 16 | cv2.imshow('m',mask2*128) 17 | cv2.waitKey() -------------------------------------------------------------------------------- /video2img.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import numpy as np 3 | import cv2 4 | import sys 5 | import os 6 | 7 | def print_process_bar(percent): 8 | cnt = 50 9 | print('['+('>'*int(percent*cnt+0.5)).ljust(cnt)+']%2d%s'%(int(percent*100),'%'), end='\r') 10 | 11 | def video2img(video_file,save_dir,stride=1,resize=1,suffix='.png'): 12 | if not os.path.exists(save_dir): 13 | os.makedirs(save_dir) 14 | print('%s not exists, create it.'%save_dir) 15 | cap = cv2.VideoCapture(video_file) 16 | frame_cnt = cap.get(cv2.CAP_PROP_FRAME_COUNT) 17 | save_idx = 0 18 | ret, frame = cap.read() 19 | r,c = frame.shape[:2] 20 | print('===============video2img=================') 21 | print('video_file:%s'%video_file) 22 | print('save_dir:%s'%save_dir) 23 | print('stride:%d'%stride) 24 | print('image size:(%dx%d)'%(c,r)+' resize:%f to (%dx%d)'%(resize,int(c*resize),int(r*resize))) 25 | 26 | print('read %d frames, start converting...'%frame_cnt) 27 | whole_cnt = frame_cnt/stride 28 | while ret: 29 | cv2.imwrite(os.path.join(save_dir,'%06d%s'%(save_idx,suffix)),cv2.resize(frame,None,fx=resize,fy=resize)) 30 | save_idx += 1 31 | for i in range(stride): 32 | ret, frame = cap.read() 33 | print_process_bar(float(save_idx)/whole_cnt) 34 | print('\nsave %d images in %s'%(save_idx,save_dir)) 35 | 36 | if __name__ == '__main__': 37 | if len(sys.argv)==2 and '-h' in sys.argv[1]: 38 | print('Usage: video2img.py video_file save_dir [stride] [resize] [suffix]') 39 | elif len(sys.argv)==3: 40 | video2img(sys.argv[1], sys.argv[2]) 41 | elif len(sys.argv)==4: 42 | video2img(sys.argv[1], sys.argv[2], int(sys.argv[3])) 43 | elif len(sys.argv)==5: 44 | video2img(sys.argv[1], sys.argv[2], int(sys.argv[3]), float(sys.argv[4])) 45 | elif len(sys.argv)==6: 46 | video2img(sys.argv[1], sys.argv[2], int(sys.argv[3]), float(sys.argv[4]), sys.argv[5]) 47 | else: 48 | print('Usage: video2img.py video_file save_dir [stride] [resize] [suffix]') --------------------------------------------------------------------------------