Spaces:
Sleeping
Sleeping
File size: 1,784 Bytes
6a6918c 146f4af 6a6918c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
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"))
|