draft_computation/endpoint.py

from .draft_computation_app import calculate_draft
import os
import cv2

from ultralytics import YOLO

model1 = YOLO("./draft_computation/models/pose.pt")
model2 = YOLO("./draft_computation/models/seg.pt")

def combine_plots(original_image, plot1, plot2):
    combined_image_1 = cv2.addWeighted(original_image, 0.7, plot1, 0.3, 0)
    final_image = cv2.addWeighted(combined_image_1, 0.7, plot2, 0.3, 0)
    return final_image

def run(img_path_or_array):

    # image_filename = "create_ocr_dataset/images/IMG_0044_01_jpg.rf.c4e4413436401ee76e86bd92e736b908.jpg"

    results1 = model1(img_path_or_array)
    results2 = model2(img_path_or_array, conf=0.15)

    try: 
        pose_results = results1[0].keypoints.data
        segment_results = results2[0].masks.xy

        draft, mid_results = calculate_draft(pose_results, segment_results, results1[0].orig_img)
        print(draft)
    except Exception as e:
        print(f"Error processing image: {e}")
        return
    
    output = {
        "draft": draft,
        "pose_results": pose_results,
        "segment_results": segment_results,
        "original_image": results1[0].orig_img,
        "pose_image_result": results1[0].plot(),
        "segment_image_result": results2[0].plot(),
        "extracted_mark_image": mid_results['extracted_mark_image']
    }
    output['final_image_result'] = combine_plots(output['original_image'],
                                                 output['pose_image_result'],
                                                 output['segment_image_result']
                                                 )
    output.update(mid_results)
    return output

if __name__ == "__main__":
    print(run("create_ocr_dataset/images/IMG_0044_01_jpg.rf.c4e4413436401ee76e86bd92e736b908.jpg"))