mirror of
https://github.com/NanjingForestryUniversity/supermachine-tobacco.git
synced 2025-11-08 14:23:55 +00:00
remove the saving function
This commit is contained in:
li.zhenye
parent
3d720d5c2b
commit
597284c3e1
107
main.py
107
main.py
@ -8,9 +8,6 @@ from models import RgbDetector, SpecDetector, ManualTree, AnonymousColorDetector
|
||||
import cv2
|
||||
|
||||
|
||||
SAVE_IMG, SAVE_NUM = False, 30
|
||||
|
||||
|
||||
def main():
|
||||
spec_detector = SpecDetector(blk_model_path=Config.blk_model_path, pixel_model_path=Config.pixel_model_path)
|
||||
rgb_detector = RgbDetector(tobacco_model_path=Config.rgb_tobacco_model_path,
|
||||
@ -23,8 +20,6 @@ def main():
|
||||
os.mkfifo(mask_fifo_path, 0o777)
|
||||
if not os.access(rgb_fifo_path, os.F_OK):
|
||||
os.mkfifo(rgb_fifo_path, 0o777)
|
||||
if SAVE_IMG:
|
||||
img_list = []
|
||||
while True:
|
||||
fd_img = os.open(img_fifo_path, os.O_RDONLY)
|
||||
fd_rgb = os.open(rgb_fifo_path, os.O_RDONLY)
|
||||
@ -55,11 +50,6 @@ def main():
|
||||
img_data = np.frombuffer(data_total, dtype=np.float32).reshape((Config.nRows, Config.nBands, -1)) \
|
||||
.transpose(0, 2, 1)
|
||||
rgb_data = np.frombuffer(rgb_data_total, dtype=np.uint8).reshape((Config.nRgbRows, Config.nRgbCols, -1))
|
||||
if SAVE_IMG:
|
||||
SAVE_NUM -= 1
|
||||
img_list.append((rgb_data, img_data))
|
||||
if SAVE_NUM <= 0:
|
||||
break
|
||||
# 光谱识别
|
||||
mask = spec_detector.predict(img_data)
|
||||
# rgb识别
|
||||
@ -71,108 +61,13 @@ def main():
|
||||
# mask_result = mask_rgb.astype(np.uint8)
|
||||
mask_result = mask_result.repeat(Config.blk_size, axis=0).repeat(Config.blk_size, axis=1).astype(np.uint8)
|
||||
t2 = time.time()
|
||||
print(f'rgb len = {len(rgb_data)}')
|
||||
|
||||
# 写出
|
||||
fd_mask = os.open(mask_fifo_path, os.O_WRONLY)
|
||||
os.write(fd_mask, mask_result.tobytes())
|
||||
os.close(fd_mask)
|
||||
t3 = time.time()
|
||||
print(f'total time is:{t3 - t1}')
|
||||
for i, img in enumerate(img_list):
|
||||
print(f"writing img {i}...")
|
||||
cv2.imwrite(f"./{i}.png", img[0][..., ::-1])
|
||||
np.save(f'./{i}.npy', img[1])
|
||||
i += 1
|
||||
|
||||
|
||||
|
||||
def save_main():
|
||||
threshold = Config.spec_size_threshold
|
||||
rgb_threshold = Config.rgb_size_threshold
|
||||
manual_tree = ManualTree(blk_model_path=Config.blk_model_path, pixel_model_path=Config.pixel_model_path)
|
||||
tobacco_detector = AnonymousColorDetector(file_path=Config.rgb_tobacco_model_path)
|
||||
background_detector = AnonymousColorDetector(file_path=Config.rgb_background_model_path)
|
||||
total_len = Config.nRows * Config.nCols * Config.nBands * 4 # float型变量, 4个字节
|
||||
total_rgb = Config.nRgbRows * Config.nRgbCols * Config.nRgbBands * 1 # int型变量
|
||||
if not os.access(img_fifo_path, os.F_OK):
|
||||
os.mkfifo(img_fifo_path, 0o777)
|
||||
if not os.access(mask_fifo_path, os.F_OK):
|
||||
os.mkfifo(mask_fifo_path, 0o777)
|
||||
if not os.access(rgb_fifo_path, os.F_OK):
|
||||
os.mkfifo(rgb_fifo_path, 0o777)
|
||||
img_list = []
|
||||
idx = 0
|
||||
while idx <= 30:
|
||||
idx += 1
|
||||
fd_img = os.open(img_fifo_path, os.O_RDONLY)
|
||||
fd_rgb = os.open(rgb_fifo_path, os.O_RDONLY)
|
||||
data = os.read(fd_img, total_len)
|
||||
|
||||
# 读取(开启一个管道)
|
||||
if len(data) < 3:
|
||||
threshold = int(float(data))
|
||||
print("[INFO] Get threshold: ", threshold)
|
||||
continue
|
||||
else:
|
||||
data_total = data
|
||||
rgb_data = os.read(fd_rgb, total_rgb)
|
||||
if len(rgb_data) < 3:
|
||||
rgb_threshold = int(float(rgb_data))
|
||||
print(rgb_threshold)
|
||||
continue
|
||||
else:
|
||||
rgb_data_total = rgb_data
|
||||
os.close(fd_img)
|
||||
os.close(fd_rgb)
|
||||
|
||||
# 识别
|
||||
t1 = time.time()
|
||||
img_data = np.frombuffer(data_total, dtype=np.float32).reshape((Config.nRows, Config.nBands, -1)). \
|
||||
transpose(0, 2, 1)
|
||||
rgb_data = np.frombuffer(rgb_data_total, dtype=np.uint8).reshape((Config.nRgbRows, Config.nRgbCols, -1))
|
||||
img_list.append((rgb_data.copy(), img_data.copy()))
|
||||
|
||||
pixel_predict_result = manual_tree.pixel_predict_ml_dilation(data=img_data, iteration=1)
|
||||
blk_predict_result = manual_tree.blk_predict(data=img_data)
|
||||
rgb_data = tobacco_detector.pretreatment(rgb_data)
|
||||
# print(rgb_data.shape)
|
||||
rgb_predict_result = 1 - (background_detector.predict(rgb_data, threshold_low=Config.threshold_low,
|
||||
threshold_high=Config.threshold_high) |
|
||||
tobacco_detector.swell(tobacco_detector.predict(rgb_data,
|
||||
threshold_low=Config.threshold_low,
|
||||
threshold_high=Config.threshold_high)))
|
||||
mask_rgb = rgb_predict_result.reshape(Config.nRows, Config.nCols // Config.blk_size, Config.blk_size) \
|
||||
.sum(axis=2).reshape(Config.nRows // 4, Config.blk_size, Config.nCols // Config.blk_size) \
|
||||
.sum(axis=1)
|
||||
mask_rgb[mask_rgb <= rgb_threshold] = 0
|
||||
mask_rgb[mask_rgb > rgb_threshold] = 1
|
||||
mask = (pixel_predict_result & blk_predict_result).astype(np.uint8)
|
||||
mask = mask.reshape(Config.nRows, Config.nCols // Config.blk_size, Config.blk_size) \
|
||||
.sum(axis=2).reshape(Config.nRows // 4, Config.blk_size, Config.nCols // Config.blk_size) \
|
||||
.sum(axis=1)
|
||||
mask[mask <= threshold] = 0
|
||||
mask[mask > threshold] = 1
|
||||
mask_result = (mask | mask_rgb).astype(np.uint8)
|
||||
# mask_result = mask_rgb
|
||||
mask_result = mask_result.repeat(Config.blk_size, axis=0).repeat(Config.blk_size, axis=1).astype(np.uint8)
|
||||
t2 = time.time()
|
||||
print(f'rgb len = {len(rgb_data)}')
|
||||
|
||||
# 写出
|
||||
fd_mask = os.open(mask_fifo_path, os.O_WRONLY)
|
||||
os.write(fd_mask, mask_result.tobytes())
|
||||
os.close(fd_mask)
|
||||
t3 = time.time()
|
||||
print(f'total time is:{t3 - t1}')
|
||||
i = 0
|
||||
print("Stop Serving")
|
||||
for img in img_list:
|
||||
print(f"writing img {i}...")
|
||||
cv2.imwrite(f"./{i}.png", img[0][..., ::-1])
|
||||
np.save(f'./{i}.npy', img[1])
|
||||
i += 1
|
||||
print("save success")
|
||||
print(f'total time is:{t3 - t1}\n')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
Loading…
Reference in New Issue
Block a user