In [116]:
import cv2
from PIL import Image, ImageDraw
import pandas as pd
import os
import matplotlib.pyplot as plt
%matplotlib inline
In [158]:
annotationts = pd.read_csv('datasets/conv/annotations/groundTruth_conv5A_test1.csv',
names=['frame', 'x1', 'y1', 'x2', 'y2', 'x3', 'y3'])
with_people = annotationts[(annotationts.x1>0)|(annotationts.x2>0)|(annotationts.x3>0)]
background = annotationts[~(annotationts.index.isin(with_people.index))]
print 'all: ', annotationts.shape[0]
print 'bgr: ', background.shape[0]
print 'persons:', with_people.shape[0]
In [182]:
directory = 'datasets/conv/conv_5A/conv5A_test1/'
image_filenames = sorted(os.listdir(directory))
draw_is = True
# background_is = True
# if background_is:
# for i in background.index:
# image_filename = image_filenames[i]
# img = Image.open(directory+image_filename)
# img.save('datasets/conv/conv_4A/test3_bgr/'+image_filename)
for i in with_people.index:
# open image
image_filename = image_filenames[i]
img = Image.open(directory+image_filename)
draw = ImageDraw.Draw(img)
x1_center, y1_center = with_people.loc[i].x1, with_people.loc[i].y1
x2_center, y2_center = with_people.loc[i].x2, with_people.loc[i].y2
x3_center, y3_center = with_people.loc[i].x3, with_people.loc[i].y3
bboxs = []
for x_center, y_center in zip([x1_center, x2_center, x3_center], [y1_center, y2_center, y3_center]):
if 50<y_center<90:
xmin, ymin, xmax, ymax = x_center-30, y_center-15, x_center+20, y_center+70
bboxs.append((xmin, ymin, xmax, ymax))
elif 90<y_center<110:
xmin, ymin, xmax, ymax = x_center-20, y_center-18, x_center+20, y_center+140
bboxs.append((xmin, ymin, xmax, ymax))
elif 110<y_center<150:
xmin, ymin, xmax, ymax = x_center-35, y_center-23, x_center+25, y_center+170
bboxs.append((xmin, ymin, xmax, ymax))
elif 150<y_center<200:
xmin, ymin, xmax, ymax = x_center-65, y_center-20, x_center+30, y_center+200
bboxs.append((xmin, ymin, xmax, ymax))
elif 200<y_center<300:
xmin, ymin, xmax, ymax = x_center-75, y_center-30, x_center+40, y_center+220
bboxs.append((xmin, ymin, xmax, ymax))
elif 330<y_center<400:
xmin, ymin, xmax, ymax = x_center-75, y_center-30, x_center+75, y_center+180
bboxs.append((xmin, ymin, xmax, ymax))
elif 400<y_center<460:
xmin, ymin, xmax, ymax = x_center-75, y_center-50, x_center+60, y_center+150
bboxs.append((xmin, ymin, xmax, ymax))
elif y_center>460:
xmin, ymin, xmax, ymax = x_center-50, y_center-50, x_center+50, y_center+40
bboxs.append((xmin, ymin, xmax, ymax))
if draw_is:
draw.rectangle((xmin, ymin, xmax, ymax))
# plt.title(str(x_center)+' '+str(y_center))
# plt.imshow(img)
# plt.show()
img.save('datasets/conv/conv_5A/test1_rect/'+image_filename)
# save annotations in pascal voc format
annotation_folder = 'datasets/conv/conv_5A/test1_annotations/'
f = open(annotation_folder + image_filename[:-4] + '.xml','w')
line = "<annotation>" + '\n'
f.write(line)
line = '\t\t<folder>' + "folder" + '</folder>' + '\n'
f.write(line)
line = '\t\t<filename>' + image_filename[:-4] + '</filename>' + '\n'
f.write(line)
line = '\t\t<source>\n\t\t<database>Source</database>\n\t</source>\n'
f.write(line)
(width, height) = img.size
line = '\t<size>\n\t\t<width>'+ str(width) + '</width>\n\t\t<height>' + str(height) + '</height>\n\t'
line += '\t<depth>3</depth>\n\t</size>'
f.write(line)
line = '\n\t<segmented>Unspecified</segmented>'
f.write(line)
for bbox in bboxs:
xmin, ymin, xmax, ymax = bbox[0], bbox[1], bbox[2], bbox[3]
line = '\n\t<object>'
line += '\n\t\t<name>person</name>\n\t\t<pose>Unspecified</pose>'
line += '\n\t\t<truncated>Unspecified</truncated>\n\t\t<difficult>0</difficult>'
line += '\n\t\t<bndbox>\n\t\t\t<xmin>' + str(xmin) + '</xmin>'
line += '\n\t\t\t<ymin>' + str(ymin) + '</ymin>'
line += '\n\t\t\t<xmax>' + str(xmax) + '</xmax>'
line += '\n\t\t\t<ymax>' + str(ymax) + '</ymax>'
line += '\n\t\t</bndbox>'
line += '\n\t</object>\n'
f.write(line)
line = "</annotation>" + '\n'
f.write(line)
f.close()
In [97]:
f = open('test.txt','w')
for i in with_people.index:
line = image_filenames[i][:-4]+'\n'
f.write(line)
f.close()
In [ ]:
# for conv4a
for x_center, y_center in zip([x1_center, x2_center, x3_center], [y1_center, y2_center, y3_center]):
if 50<y_center<90:
xmin, ymin, xmax, ymax = x_center-30, y_center-15, x_center+15, y_center+110
bboxs.append((xmin, ymin, xmax, ymax))
elif 90<y_center<110:
xmin, ymin, xmax, ymax = x_center-40, y_center-18, x_center+20, y_center+140
bboxs.append((xmin, ymin, xmax, ymax))
elif 110<y_center<150:
xmin, ymin, xmax, ymax = x_center-55, y_center-20, x_center+30, y_center+180
bboxs.append((xmin, ymin, xmax, ymax))
elif 150<y_center<200:
xmin, ymin, xmax, ymax = x_center-65, y_center-20, x_center+40, y_center+220
bboxs.append((xmin, ymin, xmax, ymax))
elif 200<y_center<300:
xmin, ymin, xmax, ymax = x_center-75, y_center-30, x_center+50, y_center+260
bboxs.append((xmin, ymin, xmax, ymax))
elif 300<y_center<400:
xmin, ymin, xmax, ymax = x_center-75, y_center-30, x_center+75, y_center+270
bboxs.append((xmin, ymin, xmax, ymax))
elif 400<y_center<500:
xmin, ymin, xmax, ymax = x_center-75, y_center-40, x_center+60, y_center+300
bboxs.append((xmin, ymin, xmax, ymax))
elif y_center>500:
xmin, ymin, xmax, ymax = x_center-50, y_center-50, x_center+50, y_center+50
bboxs.append((xmin, ymin, xmax, ymax))
if draw_is:
draw.rectangle((xmin, ymin, xmax, ymax))