-
Notifications
You must be signed in to change notification settings - Fork 0
/
MainCode.py
61 lines (47 loc) · 1.97 KB
/
MainCode.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
# -*- coding: utf-8 -*-
"""
Modified on Fri Jul 2 19:55:47 2021
Author: Vaishak
Descrption : This code demonstrates the object segmentation using a pre-trained model.
"""
## Import required modules/packages
import os
import numpy as np
from cv2 import cv2
from pixellib.semantic import semantic_segmentation
def load_images_from_folder(folder)->list:
'''
Description : Loads the images from the given directory
Input argument(s) : folder (type = str) - Input directory path to input images
Returns : images (type = list) - List of image
'''
images = []
for filename in os.listdir(folder):
img = cv2.imread(os.path.join(folder, filename))
if img is not None:
images.append(img)
return images
## Main/Driver Function
if __name__ == "__main__":
# Load Input images
IMAGEPATH = 'Input Images'
COLORIMAGELIST = load_images_from_folder(IMAGEPATH)
# Process each image
for idx, image in enumerate(COLORIMAGELIST):
# Segment image
segment_image = semantic_segmentation()
segment_image.load_pascalvoc_model("deeplabv3_xception_tf_dim_ordering_tf_kernels.h5")
process_frame = True
segmap, seg_overlay = segment_image.segmentAsPascalvoc(image, process_frame)
# Post-Processing on image
gray = cv2.cvtColor(seg_overlay, cv2.COLOR_BGR2GRAY)
ret2, mask = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
# Create a mask to the replace background pixels with the grayscale value "255"
seg_output = cv2.bitwise_and(image, image, mask=Mask)
b_plane = (seg_output[:, :, 2] == 0)
g_plane = (seg_output[:, :, 1] == 0)
r_plane = (seg_output[:, :, 0] == 0)
black = np.where(r_plane & g_plane & b_plane)
seg_output[black] = (255, 255, 255)
# Write output image
cv2.imwrite("Output images/seg_output"+ str(idx)+".jpg", seg_output)