modelscope
/
ModelScope
Not watched
1
0
Fork 0
Code
Releases 0 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: fc6d0c64bc
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'fc6d0c64bc'
${ noResults }
ModelScope/modelscope/utils/cv/image_utils.py

451 lines
16 kB
Raw Blame History 

  1. # Copyright (c) Alibaba, Inc. and its affiliates.

  2. 

  3. import cv2

  4. import matplotlib.pyplot as plt

  5. import numpy as np

  6. 

  7. from modelscope.outputs import OutputKeys

  8. from modelscope.preprocessors.image import load_image

  9. 

  10. 

  11. def numpy_to_cv2img(img_array):

  12.     """to convert a np.array with shape(h, w) to cv2 img

  13. 

  14.     Args:

  15.         img_array (np.array): input data

  16. 

  17.     Returns:

  18.         cv2 img

  19.     """

  20.     img_array = (img_array - img_array.min()) / (

  21.         img_array.max() - img_array.min() + 1e-5)

  22.     img_array = (img_array * 255).astype(np.uint8)

  23.     img_array = cv2.applyColorMap(img_array, cv2.COLORMAP_JET)

  24.     return img_array

  25. 

  26. 

  27. def draw_joints(image, np_kps, score, threshold=0.2):

  28.     lst_parent_ids_17 = [0, 0, 0, 1, 2, 0, 0, 5, 6, 7, 8, 5, 6, 11, 12, 13, 14]

  29.     lst_left_ids_17 = [1, 3, 5, 7, 9, 11, 13, 15]

  30.     lst_right_ids_17 = [2, 4, 6, 8, 10, 12, 14, 16]

  31. 

  32.     lst_parent_ids_15 = [0, 0, 1, 2, 3, 1, 5, 6, 14, 8, 9, 14, 11, 12, 1]

  33.     lst_left_ids_15 = [2, 3, 4, 8, 9, 10]

  34.     lst_right_ids_15 = [5, 6, 7, 11, 12, 13]

  35. 

  36.     if np_kps.shape[0] == 17:

  37.         lst_parent_ids = lst_parent_ids_17

  38.         lst_left_ids = lst_left_ids_17

  39.         lst_right_ids = lst_right_ids_17

  40. 

  41.     elif np_kps.shape[0] == 15:

  42.         lst_parent_ids = lst_parent_ids_15

  43.         lst_left_ids = lst_left_ids_15

  44.         lst_right_ids = lst_right_ids_15

  45. 

  46.     for i in range(len(lst_parent_ids)):

  47.         pid = lst_parent_ids[i]

  48.         if i == pid:

  49.             continue

  50. 

  51.         if (score[i] < threshold or score[1] < threshold):

  52.             continue

  53. 

  54.         if i in lst_left_ids and pid in lst_left_ids:

  55.             color = (0, 255, 0)

  56.         elif i in lst_right_ids and pid in lst_right_ids:

  57.             color = (255, 0, 0)

  58.         else:

  59.             color = (0, 255, 255)

  60. 

  61.         cv2.line(image, (int(np_kps[i, 0]), int(np_kps[i, 1])),

  62.                  (int(np_kps[pid][0]), int(np_kps[pid, 1])), color, 3)

  63. 

  64.     for i in range(np_kps.shape[0]):

  65.         if score[i] < threshold:

  66.             continue

  67.         cv2.circle(image, (int(np_kps[i, 0]), int(np_kps[i, 1])), 5,

  68.                    (0, 0, 255), -1)

  69. 

  70. 

  71. def draw_box(image, box):

  72.     cv2.rectangle(image, (int(box[0]), int(box[1])),

  73.                   (int(box[2]), int(box[3])), (0, 0, 255), 2)

  74. 

  75. 

  76. def realtime_object_detection_bbox_vis(image, bboxes):

  77.     for bbox in bboxes:

  78.         cv2.rectangle(image, (bbox[0], bbox[1]), (bbox[2], bbox[3]),

  79.                       (255, 0, 0), 2)

  80.     return image

  81. 

  82. 

  83. def draw_keypoints(output, original_image):

  84.     poses = np.array(output[OutputKeys.KEYPOINTS])

  85.     scores = np.array(output[OutputKeys.SCORES])

  86.     boxes = np.array(output[OutputKeys.BOXES])

  87.     assert len(poses) == len(scores) and len(poses) == len(boxes)

  88.     image = cv2.imread(original_image, -1)

  89.     for i in range(len(poses)):

  90.         draw_box(image, np.array(boxes[i]))

  91.         draw_joints(image, np.array(poses[i]), np.array(scores[i]))

  92.     return image

  93. 

  94. 

  95. def draw_106face_keypoints(in_path,

  96.                            keypoints,

  97.                            boxes,

  98.                            scale=4.0,

  99.                            save_path=None):

  100.     face_contour_point_index = [

  101.         0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19,

  102.         20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32

  103.     ]

  104.     left_eye_brow_point_index = [33, 34, 35, 36, 37, 38, 39, 40, 41, 33]

  105.     right_eye_brow_point_index = [42, 43, 44, 45, 46, 47, 48, 49, 50, 42]

  106.     left_eye_point_index = [66, 67, 68, 69, 70, 71, 72, 73, 66]

  107.     right_eye_point_index = [75, 76, 77, 78, 79, 80, 81, 82, 75]

  108.     nose_bridge_point_index = [51, 52, 53, 54]

  109.     nose_contour_point_index = [55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65]

  110.     mouth_outer_point_index = [

  111.         84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 84

  112.     ]

  113.     mouth_inter_point_index = [96, 97, 98, 99, 100, 101, 102, 103, 96]

  114. 

  115.     img = cv2.imread(in_path)

  116. 

  117.     for i in range(len(boxes)):

  118.         draw_box(img, np.array(boxes[i]))

  119. 

  120.     image = cv2.resize(img, dsize=None, fx=scale, fy=scale)

  121. 

  122.     def draw_line(point_index, image, point):

  123.         for i in range(len(point_index) - 1):

  124.             cur_index = point_index[i]

  125.             next_index = point_index[i + 1]

  126.             cur_pt = (int(point[cur_index][0] * scale),

  127.                       int(point[cur_index][1] * scale))

  128.             next_pt = (int(point[next_index][0] * scale),

  129.                        int(point[next_index][1] * scale))

  130.             cv2.line(image, cur_pt, next_pt, (0, 0, 255), thickness=2)

  131. 

  132.     for i in range(len(keypoints)):

  133.         points = keypoints[i]

  134. 

  135.         draw_line(face_contour_point_index, image, points)

  136.         draw_line(left_eye_brow_point_index, image, points)

  137.         draw_line(right_eye_brow_point_index, image, points)

  138.         draw_line(left_eye_point_index, image, points)

  139.         draw_line(right_eye_point_index, image, points)

  140.         draw_line(nose_bridge_point_index, image, points)

  141.         draw_line(nose_contour_point_index, image, points)

  142.         draw_line(mouth_outer_point_index, image, points)

  143.         draw_line(mouth_inter_point_index, image, points)

  144. 

  145.         size = len(points)

  146.         for i in range(size):

  147.             x = int(points[i][0])

  148.             y = int(points[i][1])

  149.             cv2.putText(image, str(i), (int(x * scale), int(y * scale)),

  150.                         cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)

  151.             cv2.circle(image, (int(x * scale), int(y * scale)), 2, (0, 255, 0),

  152.                        cv2.FILLED)

  153. 

  154.     if save_path is not None:

  155.         cv2.imwrite(save_path, image)

  156. 

  157.     return image

  158. 

  159. 

  160. def draw_face_detection_no_lm_result(img_path, detection_result):

  161.     bboxes = np.array(detection_result[OutputKeys.BOXES])

  162.     scores = np.array(detection_result[OutputKeys.SCORES])

  163.     img = cv2.imread(img_path)

  164.     assert img is not None, f"Can't read img: {img_path}"

  165.     for i in range(len(scores)):

  166.         bbox = bboxes[i].astype(np.int32)

  167.         x1, y1, x2, y2 = bbox

  168.         score = scores[i]

  169.         cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), 2)

  170.         cv2.putText(

  171.             img,

  172.             f'{score:.2f}', (x1, y2),

  173.             1,

  174.             1.0, (0, 255, 0),

  175.             thickness=1,

  176.             lineType=8)

  177.     print(f'Found {len(scores)} faces')

  178.     return img

  179. 

  180. 

  181. def draw_facial_expression_result(img_path, facial_expression_result):

  182.     scores = facial_expression_result[OutputKeys.SCORES]

  183.     labels = facial_expression_result[OutputKeys.LABELS]

  184.     label = labels[np.argmax(scores)]

  185.     img = cv2.imread(img_path)

  186.     assert img is not None, f"Can't read img: {img_path}"

  187.     cv2.putText(

  188.         img,

  189.         'facial expression: {}'.format(label), (10, 10),

  190.         1,

  191.         1.0, (0, 255, 0),

  192.         thickness=1,

  193.         lineType=8)

  194.     print('facial expression: {}'.format(label))

  195.     return img

  196. 

  197. 

  198. def draw_face_detection_result(img_path, detection_result):

  199.     bboxes = np.array(detection_result[OutputKeys.BOXES])

  200.     kpss = np.array(detection_result[OutputKeys.KEYPOINTS])

  201.     scores = np.array(detection_result[OutputKeys.SCORES])

  202.     img = cv2.imread(img_path)

  203.     assert img is not None, f"Can't read img: {img_path}"

  204.     for i in range(len(scores)):

  205.         bbox = bboxes[i].astype(np.int32)

  206.         kps = kpss[i].reshape(-1, 2).astype(np.int32)

  207.         score = scores[i]

  208.         x1, y1, x2, y2 = bbox

  209.         cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), 2)

  210.         for kp in kps:

  211.             cv2.circle(img, tuple(kp), 1, (0, 0, 255), 1)

  212.         cv2.putText(

  213.             img,

  214.             f'{score:.2f}', (x1, y2),

  215.             1,

  216.             1.0, (0, 255, 0),

  217.             thickness=1,

  218.             lineType=8)

  219.     print(f'Found {len(scores)} faces')

  220.     return img

  221. 

  222. 

  223. def draw_card_detection_result(img_path, detection_result):

  224. 

  225.     def warp_img(src_img, kps, ratio):

  226.         short_size = 500

  227.         if ratio > 1:

  228.             obj_h = short_size

  229.             obj_w = int(obj_h * ratio)

  230.         else:

  231.             obj_w = short_size

  232.             obj_h = int(obj_w / ratio)

  233.         input_pts = np.float32([kps[0], kps[1], kps[2], kps[3]])

  234.         output_pts = np.float32([[0, obj_h - 1], [0, 0], [obj_w - 1, 0],

  235.                                  [obj_w - 1, obj_h - 1]])

  236.         M = cv2.getPerspectiveTransform(input_pts, output_pts)

  237.         obj_img = cv2.warpPerspective(src_img, M, (obj_w, obj_h))

  238.         return obj_img

  239. 

  240.     bboxes = np.array(detection_result[OutputKeys.BOXES])

  241.     kpss = np.array(detection_result[OutputKeys.KEYPOINTS])

  242.     scores = np.array(detection_result[OutputKeys.SCORES])

  243.     img_list = []

  244.     ver_col = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (0, 255, 255)]

  245.     img = cv2.imread(img_path)

  246.     img_list += [img]

  247.     assert img is not None, f"Can't read img: {img_path}"

  248.     for i in range(len(scores)):

  249.         bbox = bboxes[i].astype(np.int32)

  250.         kps = kpss[i].reshape(-1, 2).astype(np.int32)

  251.         _w = (kps[0][0] - kps[3][0])**2 + (kps[0][1] - kps[3][1])**2

  252.         _h = (kps[0][0] - kps[1][0])**2 + (kps[0][1] - kps[1][1])**2

  253.         ratio = 1.59 if _w >= _h else 1 / 1.59

  254.         card_img = warp_img(img, kps, ratio)

  255.         img_list += [card_img]

  256.         score = scores[i]

  257.         x1, y1, x2, y2 = bbox

  258.         cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), 4)

  259.         for k, kp in enumerate(kps):

  260.             cv2.circle(img, tuple(kp), 1, color=ver_col[k], thickness=10)

  261.         cv2.putText(

  262.             img,

  263.             f'{score:.2f}', (x1, y2),

  264.             1,

  265.             1.0, (0, 255, 0),

  266.             thickness=1,

  267.             lineType=8)

  268.     return img_list

  269. 

  270. 

  271. def created_boxed_image(image_in, box):

  272.     image = load_image(image_in)

  273.     img = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)

  274.     cv2.rectangle(img, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])),

  275.                   (0, 255, 0), 3)

  276.     return img

  277. 

  278. 

  279. def show_video_tracking_result(video_in_path, bboxes, video_save_path):

  280.     cap = cv2.VideoCapture(video_in_path)

  281.     for i in range(len(bboxes)):

  282.         box = bboxes[i]

  283.         success, frame = cap.read()

  284.         if success is False:

  285.             raise Exception(video_in_path,

  286.                             ' can not be correctly decoded by OpenCV.')

  287.         if i == 0:

  288.             size = (frame.shape[1], frame.shape[0])

  289.             fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')

  290.             video_writer = cv2.VideoWriter(video_save_path, fourcc,

  291.                                            cap.get(cv2.CAP_PROP_FPS), size,

  292.                                            True)

  293.         cv2.rectangle(frame, (box[0], box[1]), (box[2], box[3]), (0, 255, 0),

  294.                       5)

  295.         video_writer.write(frame)

  296.     video_writer.release

  297.     cap.release()

  298. 

  299. 

  300. def show_video_object_detection_result(video_in_path, bboxes_list, labels_list,

  301.                                        video_save_path):

  302. 

  303.     PALETTE = {

  304.         'person': [128, 0, 0],

  305.         'bicycle': [128, 128, 0],

  306.         'car': [64, 0, 0],

  307.         'motorcycle': [0, 128, 128],

  308.         'bus': [64, 128, 0],

  309.         'truck': [192, 128, 0],

  310.         'traffic light': [64, 0, 128],

  311.         'stop sign': [192, 0, 128],

  312.     }

  313.     from tqdm import tqdm

  314.     import math

  315.     cap = cv2.VideoCapture(video_in_path)

  316.     with tqdm(total=len(bboxes_list)) as pbar:

  317.         pbar.set_description(

  318.             'Writing results to video: {}'.format(video_save_path))

  319.         for i in range(len(bboxes_list)):

  320.             bboxes = bboxes_list[i].astype(int)

  321.             labels = labels_list[i]

  322.             success, frame = cap.read()

  323.             if success is False:

  324.                 raise Exception(video_in_path,

  325.                                 ' can not be correctly decoded by OpenCV.')

  326.             if i == 0:

  327.                 size = (frame.shape[1], frame.shape[0])

  328.                 fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')

  329.                 video_writer = cv2.VideoWriter(video_save_path, fourcc,

  330.                                                cap.get(cv2.CAP_PROP_FPS), size,

  331.                                                True)

  332. 

  333.             FONT_SCALE = 1e-3  # Adjust for larger font size in all images

  334.             THICKNESS_SCALE = 1e-3  # Adjust for larger thickness in all images

  335.             TEXT_Y_OFFSET_SCALE = 1e-2  # Adjust for larger Y-offset of text and bounding box

  336.             H, W, _ = frame.shape

  337.             zeros_mask = np.zeros((frame.shape)).astype(np.uint8)

  338.             for bbox, l in zip(bboxes, labels):

  339.                 cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]),

  340.                               PALETTE[l], 1)

  341.                 cv2.putText(

  342.                     frame,

  343.                     l, (bbox[0], bbox[1] - int(TEXT_Y_OFFSET_SCALE * H)),

  344.                     fontFace=cv2.FONT_HERSHEY_TRIPLEX,

  345.                     fontScale=min(H, W) * FONT_SCALE,

  346.                     thickness=math.ceil(min(H, W) * THICKNESS_SCALE),

  347.                     color=PALETTE[l])

  348.                 zeros_mask = cv2.rectangle(

  349.                     zeros_mask, (bbox[0], bbox[1]), (bbox[2], bbox[3]),

  350.                     color=PALETTE[l],

  351.                     thickness=-1)

  352. 

  353.             frame = cv2.addWeighted(frame, 1., zeros_mask, .65, 0)

  354.             video_writer.write(frame)

  355.             pbar.update(1)

  356.     video_writer.release

  357.     cap.release()

  358. 

  359. 

  360. def panoptic_seg_masks_to_image(masks):

  361.     draw_img = np.zeros([masks[0].shape[0], masks[0].shape[1], 3])

  362.     from mmdet.core.visualization.palette import get_palette

  363.     mask_palette = get_palette('coco', 133)

  364. 

  365.     from mmdet.core.visualization.image import _get_bias_color

  366.     taken_colors = set([0, 0, 0])

  367.     for i, mask in enumerate(masks):

  368.         color_mask = mask_palette[i]

  369.         while tuple(color_mask) in taken_colors:

  370.             color_mask = _get_bias_color(color_mask)

  371.         taken_colors.add(tuple(color_mask))

  372. 

  373.         mask = mask.astype(bool)

  374.         draw_img[mask] = color_mask

  375. 

  376.     return draw_img

  377. 

  378. 

  379. def semantic_seg_masks_to_image(masks):

  380.     from mmdet.core.visualization.palette import get_palette

  381.     mask_palette = get_palette('coco', 133)

  382. 

  383.     draw_img = np.zeros([masks[0].shape[0], masks[0].shape[1], 3])

  384. 

  385.     for i, mask in enumerate(masks):

  386.         color_mask = mask_palette[i]

  387.         mask = mask.astype(bool)

  388.         draw_img[mask] = color_mask

  389.     return draw_img

  390. 

  391. 

  392. def show_video_summarization_result(video_in_path, result, video_save_path):

  393.     frame_indexes = result[OutputKeys.OUTPUT]

  394.     cap = cv2.VideoCapture(video_in_path)

  395.     for i in range(len(frame_indexes)):

  396.         idx = frame_indexes[i]

  397.         success, frame = cap.read()

  398.         if success is False:

  399.             raise Exception(video_in_path,

  400.                             ' can not be correctly decoded by OpenCV.')

  401.         if i == 0:

  402.             size = (frame.shape[1], frame.shape[0])

  403.             fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')

  404.             video_writer = cv2.VideoWriter(video_save_path, fourcc,

  405.                                            cap.get(cv2.CAP_PROP_FPS), size,

  406.                                            True)

  407.         if idx == 1:

  408.             video_writer.write(frame)

  409.     video_writer.release()

  410.     cap.release()

  411. 

  412. 

  413. def show_image_object_detection_auto_result(img_path,

  414.                                             detection_result,

  415.                                             save_path=None):

  416.     scores = detection_result[OutputKeys.SCORES]

  417.     labels = detection_result[OutputKeys.LABELS]

  418.     bboxes = detection_result[OutputKeys.BOXES]

  419.     img = cv2.imread(img_path)

  420.     assert img is not None, f"Can't read img: {img_path}"

  421. 

  422.     for (score, label, box) in zip(scores, labels, bboxes):

  423.         cv2.rectangle(img, (int(box[0]), int(box[1])),

  424.                       (int(box[2]), int(box[3])), (0, 0, 255), 2)

  425.         cv2.putText(

  426.             img,

  427.             f'{score:.2f}', (int(box[0]), int(box[1])),

  428.             1,

  429.             1.0, (0, 255, 0),

  430.             thickness=1,

  431.             lineType=8)

  432.         cv2.putText(

  433.             img,

  434.             label, (int((box[0] + box[2]) * 0.5), int(box[1])),

  435.             1,

  436.             1.0, (0, 255, 0),

  437.             thickness=1,

  438.             lineType=8)

  439. 

  440.     if save_path is not None:

  441.         cv2.imwrite(save_path, img)

  442.     return img

  443. 

  444. 

  445. def depth_to_color(depth):

  446.     colormap = plt.get_cmap('plasma')

  447.     depth_color = (colormap(

  448.         (depth.max() - depth) / depth.max()) * 2**8).astype(np.uint8)[:, :, :3]

  449.     depth_color = cv2.cvtColor(depth_color, cv2.COLOR_RGB2BGR)

  450.     return depth_color