line_ext.py

#!/usr/bin/env python
import sys
import os
import shutil
import math
import matplotlib.pyplot as plt

import cv2
import numpy as np

def line_ex(src_img, DEBUG):
    hsv_img = cv2.cvtColor(src_img, cv2.COLOR_BGR2HSV_FULL)             
    h, s, v = hsv_img[:,:,0], hsv_img[:,:,1], hsv_img[:,:,2]

    lower = np.array([120,18,2], dtype=np.uint8)
    upper = np.array([254,78,32], dtype=np.uint8)

    mask = cv2.inRange(hsv_img, lower, upper)
    kernel = np.array([[0, 1, 0],[1, 1, 1],[0, 1, 0]], np.uint8)
    mask = cv2.dilate(mask,kernel,iterations = 2)
    mask = cv2.erode(mask,kernel,iterations = 2)

    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

    mask = cv2.bitwise_not(mask)
    mask = cv2.fastNlMeansDenoising(mask, h=20)
    ret, mask = cv2.threshold(mask, 0, 255, cv2.THRESH_OTSU)
    _, mask = cv2.threshold(mask, 225, 255, cv2.THRESH_BINARY)
    if DEBUG == 1:
        cv2.imshow("mask", mask)
        cv2.waitKey(0)
    mask_bgr = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
        
    gray = cv2.fastNlMeansDenoising(cv2.cvtColor(src_img, cv2.COLOR_BGR2GRAY), h=20)
    dilation = cv2.dilate(gray, kernel, iterations=1)
    diff = cv2.subtract(dilation, gray)
    diff_inv = 255 - diff
    _, edge = cv2.threshold(diff_inv, 225, 255, cv2.THRESH_BINARY)
    edge = cv2.cvtColor(edge, cv2.COLOR_GRAY2BGR)
    if DEBUG == 1:
        cv2.imshow("edge", edge)
        cv2.waitKey(0)
    img = cv2.bitwise_and(mask_bgr, edge)
    if DEBUG == 1:    
        cv2.imshow("edge+mask", img)
        cv2.waitKey(0)
    return img

if __name__ == "__main__":
    IMG_FILE = sys.argv[1]
    if os.path.isfile(IMG_FILE):
        DEBUG = 1
        src_img = cv2.imread(IMG_FILE)
        result_img = line_ex(src_img, DEBUG)
        sys.exit()

    if os.path.isdir(IMG_FILE):
        DEBUG = 0
        IMG_DIR = IMG_FILE
        OUT_DIR = sys.argv[2]
        files = os.listdir(IMG_DIR)
        files = sorted(files)

        for file in files:
            src_img_path = IMG_DIR + file    
            src_img = cv2.imread(src_img_path)    
            result_img = line_ex(src_img, DEBUG)
            fname = file.split('.')[0]
            cv2.imwrite(OUT_DIR + 'line-' + fname + '.jpg', result_img)