2023-01-10 01:52:32 +08:00

195 lines
5.9 KiB
Python

import numpy as np
import cv2
BELT_Y_THRESHOLD = 390
OUTPUT_SIZE = 320
def garbage_extract_no_preprocess(img, showCompare = False):
"""
Extracts the garbage from image
img: A 4:3 image
Returns the object's image with size 320x320, black borders, scale to fit
if showCompare is True, returns the image with annotations without clipping
"""
# Flip if needed
if img.shape[0] < img.shape[1]:
img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
# Resize
img = cv2.resize(img, (390, 520), interpolation=cv2.INTER_AREA)
output = img.copy()
# Chroma Keying (blue)
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
maskKey = cv2.inRange(hsv, (90,100,130),(110,255,255))
maskKey = cv2.bitwise_not(maskKey)
output = cv2.bitwise_and(output, output, mask = maskKey)
# cv2.imshow('Output', hsv)
# cv2.waitKey(0)
# Find upper borders
gray = cv2.cvtColor(output, cv2.COLOR_BGR2GRAY)
_, bw = cv2.threshold(gray, 50, 255, cv2.THRESH_BINARY)
contours, hierarchy = cv2.findContours(bw[:BELT_Y_THRESHOLD - 10, :], cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
# Fliter small dots
fil = list(filter(lambda x: cv2.contourArea(x) > 300, contours))
cv2.drawContours(output, fil, -1, 255, 3)
# cv2.imshow('Output', output)
# cv2.waitKey(0)
x,y,w,h = cv2.boundingRect(np.vstack(fil)) if len(fil) != 0 else (0, 0, 390, 520)
# Delete left & right
output[:, :x, :] = 0
output[:, x+w: , :] = 0
# Object border
rect = [x, y, x+w, 500]
# grabCut
mask = np.zeros(img.shape[:2],np.uint8)
bgdModel = np.zeros((1,65),np.float64)
fgdModel = np.zeros((1,65),np.float64)
cv2.grabCut(img, mask, rect, bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_RECT)
mask2 = np.where((mask==2)|(mask==0),0,1).astype('uint8')
output = output*mask2[:,:,np.newaxis]
if showCompare:
output = cv2.rectangle(output, rect[:2], rect[2:], (0,255,0), 3)
return output
# Clip
# output = output[rect[1]:rect[3], rect[0]:rect[2], :]
print(rect[1], rect[3], rect[0], rect[2])
h = rect[3] - rect[1]
w = rect[2] - rect[0]
if h > w:
rect[2] = min(rect[2] + (h-w)//2, output.shape[0]-1)
rect[0] = max(rect[0] - (h-w)//2, 0)
else:
rect[3] = min(rect[3] + (w-h)//2, output.shape[1]-1)
rect[1] = max(rect[1] - (w-h)//2, 0)
print(rect[1], rect[3], rect[0], rect[2])
output = img[rect[1]:rect[3], rect[0]:rect[2], :]
# Resize
h, w, c = output.shape
scale = OUTPUT_SIZE/w if w > h else OUTPUT_SIZE/h
output = cv2.resize(output, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA)
delta_w = OUTPUT_SIZE - output.shape[1]
delta_h = OUTPUT_SIZE - output.shape[0]
top, bottom = delta_h//2, delta_h-(delta_h//2)
left, right = delta_w//2, delta_w-(delta_w//2)
color = [0, 0, 0]
output = cv2.copyMakeBorder(output, top, bottom, left, right, cv2.BORDER_CONSTANT,
value=color)
return output
def garbage_extract(img, showCompare = False):
"""
Extracts the garbage from image
img: A 4:3 image
Returns the object's image with size 320x320, black borders, scale to fit
if showCompare is True, returns the image with annotations without clipping
"""
# Flip if needed
if img.shape[0] < img.shape[1]:
img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
# Resize
img = cv2.resize(img, (390, 520), interpolation=cv2.INTER_AREA)
output = img.copy()
# Chroma Keying (blue)
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
maskKey = cv2.inRange(hsv, (90,100,130),(110,255,255))
maskKey = cv2.bitwise_not(maskKey)
output = cv2.bitwise_and(output, output, mask = maskKey)
# cv2.imshow('Output', hsv)
# cv2.waitKey(0)
# Find upper borders
gray = cv2.cvtColor(output, cv2.COLOR_BGR2GRAY)
_, bw = cv2.threshold(gray, 50, 255, cv2.THRESH_BINARY)
contours, hierarchy = cv2.findContours(bw[:BELT_Y_THRESHOLD - 10, :], cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
# Fliter small dots
fil = list(filter(lambda x: cv2.contourArea(x) > 300, contours))
# cv2.drawContours(output, fil, -1, 255, 3)
# cv2.imshow('Output', output)
# cv2.waitKey(0)
x,y,w,h = cv2.boundingRect(np.vstack(fil))
# Delete left & right
output[:, :x, :] = 0
output[:, x+w: , :] = 0
# Object border
rect = (x, y, x+w, 500)
# grabCut
mask = np.zeros(img.shape[:2],np.uint8)
bgdModel = np.zeros((1,65),np.float64)
fgdModel = np.zeros((1,65),np.float64)
cv2.grabCut(img, mask, rect, bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_RECT)
mask2 = np.where((mask==2)|(mask==0),0,1).astype('uint8')
output = output*mask2[:,:,np.newaxis]
if showCompare:
output = cv2.rectangle(output, rect[:2], rect[2:], (0,255,0), 3)
return output
# Clip
output = output[rect[1]:rect[3], rect[0]:rect[2], :]
# Resize
h, w, c = output.shape
scale = OUTPUT_SIZE/w if w > h else OUTPUT_SIZE/h
output = cv2.resize(output, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA)
delta_w = OUTPUT_SIZE - output.shape[1]
delta_h = OUTPUT_SIZE - output.shape[0]
top, bottom = delta_h//2, delta_h-(delta_h//2)
left, right = delta_w//2, delta_w-(delta_w//2)
color = [0, 0, 0]
output = cv2.copyMakeBorder(output, top, bottom, left, right, cv2.BORDER_CONSTANT,
value=color)
return output
if __name__ == "__main__":
img = cv2.imread('Camera/paper/paper_leafyellow_14.jpg')
# compare = garbage_extract(img, True)
output = garbage_extract(img)
# Flip if needed
if img.shape[0] < img.shape[1]:
img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
# Resize
img = cv2.resize(img, (390, 520), interpolation=cv2.INTER_AREA)
# cv2.imshow('Compare', np.hstack([compare, img]))
cv2.imshow('Output', output)
cv2.waitKey(0)
cv2.destroyAllWindows()