try: import detectron2 except: import os os.system('pip install git+https://github.com/haya-alwarthan/detectron2.git') try: import roboflow except: os.system('pip install git+https://github.com/roboflow-ai/roboflow-python.git') import cv2 import gradio as gr import numpy as np import torch import requests import roboflow from detectron2 import model_zoo from detectron2.engine import DefaultPredictor from detectron2.config import get_cfg from detectron2.utils.visualizer import Visualizer from detectron2.data import MetadataCatalog from detectron2.utils.visualizer import ColorMode from matplotlib import patches import matplotlib.pyplot as plt from matplotlib.colors import to_rgba import numpy as np import cv2 import io #Define the path to the pretrained weights of the model keypoint_model_path= "https://huggingface.co/yahaoh/ddh-maskrcnn/resolve/main/keypoint_rcnn_binarymask.pth" #configure roboflow project rf = roboflow.Roboflow(api_key="0uv4bY5n7Vluj0yRHOOm") project = rf.workspace().project("ddh-seg") model = project.version(1).model #convert detectron format predictions to normal pairs def convert_to_pairs(three): pairs=[element for i,element in enumerate(three) if i%3!=2 ] return pairs #function to output binary mask img def segm_imf(prediction,rec_img): w= rec_img.shape[0] h=rec_img.shape[1] colors={"LOWER_BONE":"#fabee6","UPPER_BONE":"#96e7e6","MIDDLE_BONE":"#fffa5b"} # colors={"LOWER_BONE":(253,226,243),"UPPER_BONE":(173,228,219),"MIDDLE_BONE": (253,247,195)} figure, axes = plt.subplots(figsize =(h/100.0,w/100.0)) for prediction in prediction["predictions"]: points = [[p["x"], p["y"]] for p in prediction["points"]] polygon = patches.Polygon( points, linewidth=2, edgecolor=colors[prediction["class"]],facecolor= to_rgba(colors[prediction["class"]],0.4) ) axes.add_patch(polygon) plt.imshow(rec_img) plt.axis("off") plt.tight_layout() canvas = plt.get_current_fig_manager().canvas canvas.draw() buf = io.BytesIO() canvas.print_png(buf) buf.seek(0) img_arr = np.frombuffer(buf.getvalue(), dtype='uint8') img = cv2.imdecode(img_arr, cv2.IMREAD_UNCHANGED) gray_image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) return gray_image def binary_cv2(rec_img,prediction): w= rec_img.shape[0] h=rec_img.shape[1] mask = np.zeros((w, h), dtype = np.float64) for prediction in prediction["predictions"]: points = [[p["x"], p["y"]] for p in prediction["points"]] cv2.fillPoly(mask, np.array([points]).astype(np.int32), color=(255, 0, 0)) masked_gray = cv2.merge((mask,mask,mask)) return masked_gray def extend_line(p1, p2, distance=10000): diff = np.arctan2(p1[1] - p2[1], p1[0] - p2[0]) p3_x = int(p1[0] + distance*np.cos(diff)) p3_y = int(p1[1] + distance*np.sin(diff)) p4_x = int(p1[0] - distance*np.cos(diff)) p4_y = int(p1[1] - distance*np.sin(diff)) return ((p3_x, p3_y), (p4_x, p4_y)) def visualize(image,pred): op_img=image.copy() pred=[int(i) for i in pred] (p1_left,p2_left)=extend_line((pred[6],pred[7]),(pred[4],pred[5])) (p1_h,p2_h)=extend_line((pred[4],pred[5]),(pred[2],pred[3])) (p1_right,p2_right)=extend_line((pred[2],pred[3]),(pred[0],pred[1])) op_img= cv2.line(op_img, p1_left, (pred[4],pred[5]),(152, 216, 170),1) op_img= cv2.line(op_img, p1_h, p2_h,(152, 216, 170),1) op_img= cv2.line(op_img, (pred[2],pred[3]), p2_right,(152, 216, 170),1) for i in range(0,7,2): op_img = cv2.circle(op_img, (round(pred[i]),round(pred[i+1])), int ((image.shape[0]+image.shape[1])/150), (255, 150, 128), -1) return op_img #Keypoint RCNN MODEL AND META DATA SETUP KEYPOINT_NAMES = ["RU","RD","LD","LU"] KEYPOINT_FLIP_MAP = [ ("RU", "LU"), ("RD", "LD"), ] KEYPOINT_CONNECTION_RULES = [ ("RU", "RD", (102, 204, 255)), ("RD", "LD", (51, 153, 255)), ("LU", "LD", (102, 0, 204)), ] kp_meta=MetadataCatalog.get("ddh_jordan_kp_train") kp_meta.set(keypoint_names =KEYPOINT_NAMES) kp_meta.set(keypoint_flip_map =KEYPOINT_FLIP_MAP) kp_meta.set(keypoint_connection_rules =KEYPOINT_CONNECTION_RULES) kp_cfg = get_cfg() kp_cfg.merge_from_file(model_zoo.get_config_file("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml")) kp_cfg.DATALOADER.NUM_WORKERS = 2 kp_cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml") # Let training initialize from model zoo kp_cfg.MODEL.ROI_KEYPOINT_HEAD.NUM_KEYPOINTS=4 kp_cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # only has one class (landmarks). (see https://detectron2.readthedocs.io/tutorials/datasets.html#update-the-config-for-new-datasets) kp_cfg.TEST.DETECTIONS_PER_IMAGE=1 kp_cfg.MODEL.WEIGHTS = keypoint_model_path # kp_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set a custom testing threshold #check whether gpu is available if not torch.cuda.is_available(): kp_cfg.MODEL.DEVICE = "cpu" #Define the predictor of the model kp_predictor = DefaultPredictor(kp_cfg) def output_json(keypoints): landmarks_labels= ["RU","RD","LD","LU"] output= {} for i in range(len(landmarks_labels)): output[landmarks_labels[i].lower()]={'x':keypoints[2*i].item(), 'y':keypoints[2*i+1].item()} return output #Define a function to infernece from image def predict_fn(img_path): #Read and tranform input image preds=model.predict(img_path).json() og_img=cv2.imread(img_path) img=binary_cv2(og_img,preds) img= np.array(img,dtype=np.uint8) outputs = kp_predictor(img) # format is documented at https://detectron2.readthedocs.io/tutorials/models.html#model-output-format print("outputs==".format(outputs["instances"].to("cpu"))) p=np.asarray(outputs["instances"].to("cpu").pred_keypoints, dtype='float32') landmarks={} if p.size >0: p=p[0].reshape(-1) pairss=convert_to_pairs(p) landmarks=output_json(pairss) segm=segm_imf(preds,og_img) img= visualize(segm,pairss) return (img,landmarks) inputs_image = [ gr.components.Image(type="filepath", label="Upload an XRay Image of the Pelvis"), ] outputs_image = [ gr.components.Image(type="numpy", label="Output Image") ] outputs_landmarks = [ gr.components.JSON( label="Output Landmarks") ] outputs=[ gr.components.Image(type="numpy", label="Output Image"), gr.components.JSON( label="Output Landmarks") ] gr.Interface( predict_fn, inputs=inputs_image, outputs=outputs, title="Coordinates of the Landmarks", ).launch()