MLVLab/Human_Object_Interaction

1from __future__ import absolute_import2from __future__ import division3from __future__ import print_function4 5import numpy as np6import cv27 8 9vcoco_action_string = {2: 'hold', 3: 'stand', 4: 'sit', 5: 'ride', 6: 'walk',\10                       7: 'look', 8: 'hit_inst', 9: 'hit_obj', 10: 'eat_obj', \11                       11: 'eat_inst', 12: 'jump', 13: 'lay', 14: 'talk', 15: \12                       'carry', 16: 'throw', 17: 'catch', 18: 'cut_inst', 19:'cut_obj', \13                       20: 'run', 21: 'work_on_comp', 22: 'ski', 23: 'surf', 24: 'skateboard', \14                       25: 'smile', 26: 'drink', 27: 'kick', 28: 'point', 29: 'read', 30: 'snowboard'}15def draw_box_on_img(box, img,color=None):16 17    vis_img = img.copy()18    box = [int(x) for x in box]19    cv2.rectangle(vis_img, (box[0], box[1]), (box[2], box[3]), color, 2)20    draw_point=[int((box[0]+box[2])*1.0/2),int((box[1]+box[3])*1.0/2)]21    22    return vis_img,color23 24 25def draw_line_on_img_vcoco(box,line, img, class_index,color):26    27    vis_img = img.copy()28    font=cv2.FONT_HERSHEY_SIMPLEX29    x=int(box[0])+230    y=int(box[1])+231    f=int(box[1])+232    for i in range(len(class_index)):33     34        font_scale=135        font_thickness=236    37        text_size, _ = cv2.getTextSize(vcoco_action_string[class_index[i]] , font, font_scale, font_thickness)38        vis_img=cv2.rectangle(vis_img,(x,y),(x+text_size[0],y+text_size[1]+5),color[1],-1)39        40 41        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)42        y=y+text_size[1]+543   44    return vis_img,y45 46 47def draw_img_vcoco(img, output_i, top_k,threshold,color):48    list_action = []49    for action in output_i['hoi_prediction']:50        subject_id = action['subject_id']51        object_id = action['object_id']52        category_id = action['category_id']53        score = action['score']54        single_out = [subject_id,object_id,category_id,score]55        list_action.append(single_out)56    list_action = sorted(list_action, key=lambda x:x[-1], reverse=True)57    action_dict = []58    action_cate = []59    action_color=[]60    subj_box=[]61    sb={}62    sbj=[]63    for action in list_action[:top_k]:64 65        subject_id,object_id,category_id,score = action66        if score<threshold:67            break68        subject_obj = output_i['predictions'][subject_id]69        subject_box = subject_obj['bbox']70        object_obj = output_i['predictions'][object_id]71        object_box = object_obj['bbox']72 73        point_1 = [int((subject_box[0]+subject_box[2])*1.0/2),int((subject_box[1]+subject_box[3])*1.0/2)]74        point_2 = [int((object_box[0]+object_box[2])*1.0/2),int((object_box[1]+object_box[3])*1.0/2)]75  76        if [point_1,point_2] not in action_dict:77            78            img,color_hum = draw_box_on_img(subject_box, img, color[subject_obj['category_id']]['color'])79            80            img,color_obj = draw_box_on_img(object_box, img, color[object_obj['category_id']]['color'])81            82            action_dict.append([point_1,point_2])83            action_color.append([color_hum,color_obj])84            subj_box.append([int(subject_box[0]),int(subject_box[1])])    85           86        action_cate.append([])87        action_cate[action_dict.index([point_1,point_2])].append(category_id)88 89    for i,(action_item,clr) in enumerate(zip(action_dict,action_color)):90        91        img,offset = draw_line_on_img_vcoco(subj_box[i],action_item,img,action_cate[action_dict.index(action_item)],clr)92 93        for p in range(i+1,len(subj_box)):94            if subj_box[p]==subj_box[i]:95                subj_box[p][1]=offset96    return img