diff --git a/extract_information.py b/extract_information.py new file mode 100644 index 0000000..1a7a51f --- /dev/null +++ b/extract_information.py @@ -0,0 +1,63 @@ +import cv2 +import easyocr +import numpy as np + +def extract_information_from_image(img, coordinates_dict, **kwargs): + + def preprocess_image(image): + image = cv2.medianBlur(image, 3) + image = image * 0.7 + image = np.clip(image, 0, 255) + image = image.astype(np.uint8) + alpha = 1.5 + beta = 0 + image = np.clip(alpha * image + beta, 0, 255).astype(np.uint8) + image = cv2.medianBlur(image, 5) + image = cv2.GaussianBlur(image, (3, 3), 0) + return image + + + allowed_list = " 0123456789(){}[]\'\"\\/.,_" + def check(input:str): + for letter in input: + if not (letter in allowed_list): + # print('check list: ', letter) + return '--' + return input + + + if kwargs.get("reader") is None: + reader = easyocr.Reader(["en"]) + else: + reader = kwargs.get("reader") + + img = preprocess_image(img) + img = cv2.resize(img, (1280, 720)) + #cv2.imwrite("saved_test_3.jpg", img=img) + + #print(coordinates_dict) + + res = {} + for key, value in coordinates_dict.items(): + temp = reader.readtext(img[value[1]:value[3], value[0]:value[2]], detail=0) + print(temp) + try: + res[key] = ( + check(input=temp[0]) + if len(temp) >=1 else '--' + ) + except: + res[key] = '--' + return res + +if __name__ == "__main__": + img = cv2.imread("ex.jpeg") + print(extract_information_from_image(img, { + "map": [141, 378, 272, 447], + "heart_rate": [71, 80, 415, 242], + "spo2": [93, 469, 424, 605], + "sbp": [18, 310, 195, 405], + "rr": [303, 646, 452, 726], + "dbp": [257, 317, 442, 398] + })) + diff --git a/pipeline-1-single-image.py b/pipeline-1-single-image.py new file mode 100644 index 0000000..1a45016 --- /dev/null +++ b/pipeline-1-single-image.py @@ -0,0 +1,45 @@ +from ultralytics import YOLO +from classifier_model import predict_monitor +from segmentation2polygon import segmentation2polygon, do_perspective_transformation +import cv2 +from classifier_model import predict_monitor +import json +from extract_information import extract_information_from_image +import argparse + +with open('data.json', 'r') as f: + content = json.load(f) +parser = argparse.ArgumentParser(description="CloudPhysician") +parser.add_argument('-i', '--img', help='Path to image') + +def extract_info(image): + yolo_model = YOLO("./model_yolo/best.pt") + yolo_result = yolo_model.predict(image) + + yolo_mask = yolo_result[0].masks.xy + + approx_polygon = segmentation2polygon(yolo_mask[0]).reshape((-1, 2)) + + perspective_image = do_perspective_transformation(image, approx_polygon) + + class_number = predict_monitor(perspective_image) + if class_number == "1": + coordinates_dict = content['first'] + elif class_number =="2": + coordinates_dict = content['second'] + elif class_number =="3": + coordinates_dict = content['third'] + elif class_number =="4": + coordinates_dict = content['fourth'] + else: + exit(0) + return extract_information_from_image(image, coordinates_dict) + + + +if __name__ == '__main__': + args=parser.parse_args() + img = cv2.imread(args.img) + print(extract_info(img)) + + diff --git a/segmentation2polygon.py b/segmentation2polygon.py index 9106f46..4f63737 100644 --- a/segmentation2polygon.py +++ b/segmentation2polygon.py @@ -77,7 +77,7 @@ def pipeline(): segmentation = get_segmentation() segmentation = convert_str_to_array(segmentation) img_size = get_height_and_width_of_img( - "./yolo/yolo_predictions/medicakolkata_rcu_mon--8_2022_5_18_17_5_38.jpeg" + "ex.jpeg" ) polygon = from_normalised_to_pixel(img_size, segmentation) polygon = segmentation2polygon(polygon) @@ -92,7 +92,7 @@ def draw_points(polygon) -> None: Output: None """ img = cv2.imread( - "./yolo/yolo_predictions/medicakolkata_rcu_mon--8_2022_5_18_17_5_38.jpeg" + "ex.jpeg" ) for i in polygon: cv2.circle(img, (i[0], i[1]), 5, (0, 0, 255), -1) @@ -129,15 +129,22 @@ def do_perspective_transformation(image, input_array): if __name__ == "__main__": - approx_polygon = pipeline().reshape((-1, 2)) + '''approx_polygon = np.array([[ 42 , 12], + [ 72 , 692], + [1252 , 684], + [1236 , 10]]) + print(len(approx_polygon)) draw_points(approx_polygon) perspective = do_perspective_transformation( cv2.imread( - "./yolo/yolo_predictions/medicakolkata_rcu_mon--8_2022_5_18_17_5_38.jpeg" + "ex.jpeg" ), approx_polygon, ) cv2.imshow("perspective", perspective) + #cv2.imwrite('ex.jpeg', perspective) cv2.waitKey(0) - cv2.destroyAllWindows() + cv2.destroyAllWindows()''' + + print(segmentation2polygon(get_segmentation())) diff --git a/yolo.py b/yolo.py new file mode 100644 index 0000000..d58cd3b --- /dev/null +++ b/yolo.py @@ -0,0 +1,21 @@ +""" + TODO: + 1. get the .pt file of the yolo model and try to make it in such a way that it store output in our desired format +""" +from ultralytics import YOLO +from segmentation2polygon import segmentation2polygon, draw_points + +model = YOLO("./model_yolo/best.pt") +results = model.predict("ex.jpeg") + +mask = results[0].masks + +mask = mask.xy +print(mask[0]) +segmentation = segmentation2polygon(mask[0]) +print(segmentation.reshape((-1, 2))) +#with open('mask.txt', 'w') as f: + #f.write(mask[0]) +draw_points(segmentation.reshape((-1, 2))) + +