Spaces:
Runtime error
Runtime error
| from __future__ import absolute_import | |
| from __future__ import division | |
| from __future__ import print_function | |
| import numpy as np | |
| import cv2 | |
| vcoco_action_string = {2: 'hold', 3: 'stand', 4: 'sit', 5: 'ride', 6: 'walk',\ | |
| 7: 'look', 8: 'hit_inst', 9: 'hit_obj', 10: 'eat_obj', \ | |
| 11: 'eat_inst', 12: 'jump', 13: 'lay', 14: 'talk', 15: \ | |
| 'carry', 16: 'throw', 17: 'catch', 18: 'cut_inst', 19:'cut_obj', \ | |
| 20: 'run', 21: 'work_on_comp', 22: 'ski', 23: 'surf', 24: 'skateboard', \ | |
| 25: 'smile', 26: 'drink', 27: 'kick', 28: 'point', 29: 'read', 30: 'snowboard'} | |
| def draw_box_on_img(box, img,color=None): | |
| vis_img = img.copy() | |
| box = [int(x) for x in box] | |
| cv2.rectangle(vis_img, (box[0], box[1]), (box[2], box[3]), color, 2) | |
| draw_point=[int((box[0]+box[2])*1.0/2),int((box[1]+box[3])*1.0/2)] | |
| return vis_img,color | |
| def draw_line_on_img_vcoco(box,line, img, class_index,color): | |
| vis_img = img.copy() | |
| font=cv2.FONT_HERSHEY_SIMPLEX | |
| x=int(box[0])+2 | |
| y=int(box[1])+2 | |
| f=int(box[1])+2 | |
| for i in range(len(class_index)): | |
| font_scale=1 | |
| font_thickness=2 | |
| text_size, _ = cv2.getTextSize(vcoco_action_string[class_index[i]] , font, font_scale, font_thickness) | |
| vis_img=cv2.rectangle(vis_img,(x,y),(x+text_size[0],y+text_size[1]+5),color[1],-1) | |
| vis_img=cv2.putText(vis_img, vcoco_action_string[class_index[i]] ,(x,y + text_size[1] ),font,font_scale,[51,255,153],font_thickness) | |
| y=y+text_size[1]+5 | |
| return vis_img,y | |
| def draw_img_vcoco(img, output_i, top_k,threshold,color): | |
| list_action = [] | |
| for action in output_i['hoi_prediction']: | |
| subject_id = action['subject_id'] | |
| object_id = action['object_id'] | |
| category_id = action['category_id'] | |
| score = action['score'] | |
| single_out = [subject_id,object_id,category_id,score] | |
| list_action.append(single_out) | |
| list_action = sorted(list_action, key=lambda x:x[-1], reverse=True) | |
| action_dict = [] | |
| action_cate = [] | |
| action_color=[] | |
| subj_box=[] | |
| sb={} | |
| sbj=[] | |
| for action in list_action[:top_k]: | |
| subject_id,object_id,category_id,score = action | |
| if score<threshold: | |
| break | |
| subject_obj = output_i['predictions'][subject_id] | |
| subject_box = subject_obj['bbox'] | |
| object_obj = output_i['predictions'][object_id] | |
| object_box = object_obj['bbox'] | |
| point_1 = [int((subject_box[0]+subject_box[2])*1.0/2),int((subject_box[1]+subject_box[3])*1.0/2)] | |
| point_2 = [int((object_box[0]+object_box[2])*1.0/2),int((object_box[1]+object_box[3])*1.0/2)] | |
| if [point_1,point_2] not in action_dict: | |
| img,color_hum = draw_box_on_img(subject_box, img, color[subject_obj['category_id']]['color']) | |
| img,color_obj = draw_box_on_img(object_box, img, color[object_obj['category_id']]['color']) | |
| action_dict.append([point_1,point_2]) | |
| action_color.append([color_hum,color_obj]) | |
| subj_box.append([int(subject_box[0]),int(subject_box[1])]) | |
| action_cate.append([]) | |
| action_cate[action_dict.index([point_1,point_2])].append(category_id) | |
| for i,(action_item,clr) in enumerate(zip(action_dict,action_color)): | |
| img,offset = draw_line_on_img_vcoco(subj_box[i],action_item,img,action_cate[action_dict.index(action_item)],clr) | |
| for p in range(i+1,len(subj_box)): | |
| if subj_box[p]==subj_box[i]: | |
| subj_box[p][1]=offset | |
| return img |