Spaces:
Runtime error
Runtime error
File size: 3,680 Bytes
5219368 |
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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
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 |