diff --git a/research/roi.py b/research/roi.py index 580d292..0a4a8e7 100755 --- a/research/roi.py +++ b/research/roi.py @@ -5,6 +5,7 @@ import os import requests import json import subprocess +import cv2 class RoiResearch(object): def __init__(self, token=None): @@ -37,13 +38,6 @@ class RoiResearch(object): return yield sd -def parse_args(): - parser = argparse.ArgumentParser() - parser.add_argument("--token", '-t', default='3ebd8c33-f46e-4b06-bda8-4c0f5f5eb530', type=str) - parser.add_argument("--username", "-u", type=str) - parser.add_argument("--password", "-p", type=str) - return parser.parse_args() - def get_last_id(list_file): ret = None if not os.path.exists(list_file): @@ -63,9 +57,8 @@ def get_all_ids(list_file): for line in f: ret.append(json.loads(line)['id']) return ret - -def main(): - args = parse_args() + +def get_roi_data(args): rr = RoiResearch(args.token) if args.username and args.password: print(rr.login(args.username, args.password)) @@ -93,6 +86,120 @@ def main(): cmd = f'../scripts/emcli get-scan-data {id} -o {outd}' print(cmd) subprocess.call(cmd, shell=True) + +def preprocess(img): + gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) + min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(gray) + histo = cv2.calcHist([gray], [0], None, [256], [0, 256]) + total = gray.shape[0] * gray.shape[1] + s = 0 + low_thres = -1 + high_thres = -1 + low_proportion = 70 + high_proportion = 80 + for x in range(256): + i = histo[x] + s += i + ratio = s / total * 100 + if low_thres == -1 and ratio > low_proportion: + low_thres = x + if high_thres == -1 and ratio > high_proportion: + high_thres = x + if low_thres == -1: + raise Exception("no threshold found") + print(low_thres, high_thres) + for i in range(gray.shape[0]): + for j in range(gray.shape[1]): + if gray[i, j] < low_thres: + gray[i, j] = 0 + elif gray[i, j] > high_thres: + gray[i, j] = 0 + else: + gray[i, j] = 255 + return gray + +window_y = 0 +def show_img(img, name): + global window_y + cv2.imshow(name, img) + cv2.moveWindow(name, 640, window_y) + window_y += 200 + +def find_point(roi_processed, i, j): + tolerrance = 2 + for di in range(-tolerrance, tolerrance + 1): + for dj in range(-tolerrance, tolerrance + 1): + x = i + di + y = j + dj + if x < 0 or x >= roi_processed.shape[0] or y < 0 or y >= roi_processed.shape[1]: + continue + if roi_processed[x, y] > 0: + return True + +def calc_sim(frame_processed, roi_processed): + s = 0 + total = 0 + width = frame_processed.shape[1] + height = frame_processed.shape[0] + for i in range(height): + for j in range(width): + roi_val = roi_processed[i, j] + frame_val = frame_processed[i, j] + if roi_val: + total += 1 + if find_point(frame_processed, i, j): + s += 1 + if frame_val: + total += 1 + if find_point(roi_processed, i, j): + s += 1 + return s / total + +def roi_sim(frame_roi_img, roi_img): + if frame_roi_img.shape != roi_img.shape: + raise Exception("size mismatch") + frame_processed = preprocess(frame_roi_img) + roi_processed = preprocess(roi_img) + siml = calc_sim(frame_processed, roi_processed) + if False: + show_img(frame_processed, "frame_processed") + show_img(roi_processed, "roi_processed") + show_img(frame_roi_img, "frame_roi_img") + show_img(roi_img, "roi_img") + cv2.waitKey(0) + cv2.destroyAllWindows() + return siml + +def process_roi_data(id): + rd = f"data/roi/samples/{id}" + frame_file = os.path.abspath(f"{rd}/{id}-frame.jpg") + roi_file = os.path.abspath(f"{rd}/{id}-roi.jpg") + roi_img = cv2.imread(roi_file) + cmd = f'./qrtool frame_roi {frame_file}' + subprocess.check_call(cmd, shell=True, cwd=os.path.dirname(os.path.abspath(__file__)) + "/../alg") + frame_roi_img = cv2.imread(frame_file + ".roi.jpg") + size = [128, 128] + frame_roi_img = cv2.resize(frame_roi_img, size) + roi_img = cv2.resize(roi_img, size) + sim = roi_sim(frame_roi_img, roi_img) + print(f"{id}: {sim}") + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--token", '-t', default='3ebd8c33-f46e-4b06-bda8-4c0f5f5eb530', type=str) + parser.add_argument("--username", "-u", type=str) + parser.add_argument("--password", "-p", type=str) + parser.add_argument("--download", "-d", action='store_true') + parser.add_argument("--id", "-i", type=int, action='append') + return parser.parse_args() + +def main(): + args = parse_args() + if args.download: + get_roi_data(args) + if args.id: + for id in args.id: + process_roi_data(id) if __name__ == "__main__": main()