generate_cache_caltech.py

import os
import pickle
import numpy as np
import matplotlib.pyplot as plt

root_dir = 'data/caltech/train_3'
# root_dir = 'data/caltech/test'
all_img_path = os.path.join(root_dir, 'images')
all_anno_path = os.path.join(root_dir, 'annotations_new/')
res_path_gt = 'data/cache/caltech/train_gt'
res_path_nogt = 'data/cache/caltech/train_nogt'

rows, cols = 480, 640
image_data_gt, image_data_nogt = [], []

valid_count = 0
iggt_count = 0
box_count = 0
files = sorted(os.listdir(all_anno_path))
for l in range(len(files)):
    gtname = files[l]
    imgname = files[l].split('.')[0] + '.jpg'
    img_path = os.path.join(all_img_path, imgname)
    gt_path = os.path.join(all_anno_path, gtname)

    boxes = []
    ig_boxes = []
    with open(gt_path, 'rb') as fid:
        lines = fid.readlines()
    if len(lines) > 1:
        for i in range(1, len(lines)):
            info = lines[i].strip().split(' ')
            label = info[0]
            occ, ignore = info[5], info[10]
            x1, y1 = max(int(float(info[1])), 0), max(int(float(info[2])), 0)
            w, h = min(int(float(info[3])), cols - x1 - 1), min(int(float(info[4])), rows - y1 - 1)
            box = np.array([int(x1), int(y1), int(x1) + int(w), int(y1) + int(h)])
            if int(ignore) == 0:
                boxes.append(box)
            else:
                ig_boxes.append(box)
    boxes = np.array(boxes)
    ig_boxes = np.array(ig_boxes)

    annotation = {}
    annotation['filepath'] = img_path
    box_count += len(boxes)
    iggt_count += len(ig_boxes)
    annotation['bboxes'] = boxes
    annotation['ignoreareas'] = ig_boxes
    if len(boxes) == 0:
        image_data_nogt.append(annotation)
    else:
        image_data_gt.append(annotation)
        valid_count += 1
print('{} images and {} valid images, {} valid gt and {} ignored gt'.format(len(files), valid_count, box_count,
                                                                            iggt_count))

if not os.path.exists(res_path_gt):
    with open(res_path_gt, 'wb') as fid:
        pickle.dump(image_data_gt, fid, pickle.HIGHEST_PROTOCOL)
if not os.path.exists(res_path_nogt):
    with open(res_path_nogt, 'wb') as fid:
        pickle.dump(image_data_nogt, fid, pickle.HIGHEST_PROTOCOL)