RE-OWOD/datasets/save_selective_search.py

43 lines
1.6 KiB
Python

import cv2
import sys
# from drawBoxes import draw_boxes
import time
import torch
import matplotlib.pyplot as plt
from selective_search import selective_search
import numpy as np
t1_train_file_path = "/home/selective_search/all_task_test.txt"
with open(t1_train_file_path, 'r') as t1_train_file:
t1_train_list = t1_train_file.read().splitlines()
for image_num,image_name in enumerate(t1_train_list):
image_name = t1_train_list[image_num]
image_path = "/home//datasets/VOC2007/JPEGImages/" + image_name + ".jpg"
img = cv2.imread(image_path)[:, :, ::-1]
img_rgb = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
try:
boxes = selective_search(img_rgb, mode='single', random_sort=False)
new_flag = True
for boxes_i in boxes:
if new_flag:
boxes_sum = np.array(boxes_i).reshape(1,4)
new_flag = False
else:
boxes_sum = np.r_[boxes_sum,np.array(boxes_i).reshape(1,4)]
boxes_draw = boxes_sum[:50,:]
img_save = {'image_size':img.shape[0:2],'obj_boxes':boxes_draw}
except:
print("cannot compute score:",image_name)
boxes_draw = []
img_save = {'image_size':img.shape[0:2],'obj_boxes':boxes_draw}
error_save_path = "/home/selective_search_save/selective_search_test_error/" + image_name + ".pickle"
torch.save(img_save, error_save_path)
img_save_path = "/home/selective_search_save/selective_search_test/" + image_name + ".pickle"
torch.save(img_save, img_save_path)
print("image_num:",image_num)