Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python3 | |
| # -*- coding: utf-8 -*- | |
| """ | |
| Created on Tue Oct 19 19:14:47 2021 | |
| @author: dinesh | |
| """ | |
| import glob | |
| from .utils import bb_intersection_over_union_unoccluded | |
| import numpy as np | |
| from PIL import Image | |
| import datetime | |
| import cv2 | |
| import os | |
| from tqdm import tqdm | |
| def get_image(time, folder): | |
| for week_loop in range(5): | |
| try: | |
| image = np.array(Image.open(folder+'/week' +str(week_loop)+'/'+ str(time).replace(' ','T').replace(':','-').split('+')[0] + '.jpg')) | |
| break | |
| except: | |
| continue | |
| if image is None: | |
| print('file not found') | |
| return image | |
| def get_mask(segm, image): | |
| poly = np.array(segm).reshape((int(len(segm)/2), 2)) | |
| mask = image.copy()*0 | |
| cv2.fillConvexPoly(mask, poly, (255, 255, 255)) | |
| return mask | |
| def get_unoccluded(indices, tracks_all): | |
| unoccluded_indexes = [] | |
| unoccluded_index_all =[] | |
| while 1: | |
| unoccluded_clusters = [] | |
| len_unocc = len(unoccluded_indexes) | |
| for ind in indices: | |
| if ind in unoccluded_indexes: | |
| continue | |
| occ = False | |
| for ind_compare in indices: | |
| if ind_compare in unoccluded_indexes: | |
| continue | |
| if bb_intersection_over_union_unoccluded(tracks_all[ind], tracks_all[ind_compare]) > 0.01 and ind_compare != ind: | |
| occ = True | |
| if occ==False: | |
| unoccluded_indexes.extend([ind]) | |
| unoccluded_clusters.extend([ind]) | |
| if len(unoccluded_indexes) == len_unocc and len_unocc != 0: | |
| for ind in indices: | |
| if ind not in unoccluded_indexes: | |
| unoccluded_indexes.extend([ind]) | |
| unoccluded_clusters.extend([ind]) | |
| unoccluded_index_all.append(unoccluded_clusters) | |
| if len(unoccluded_indexes) > len(indices)-5: | |
| break | |
| return unoccluded_index_all | |
| def primes(n): # simple sieve of multiples | |
| odds = range(3, n+1, 2) | |
| sieve = set(sum([list(range(q*q, n+1, q+q)) for q in odds], [])) | |
| return [2] + [p for p in odds if p not in sieve] | |
| def save_image(image_read, save_path, data, path): | |
| tracks = data['tracks_all_unoccluded'] | |
| segmentations = data['segmentation_all_unoccluded'] | |
| timestamps = data['timestamps_final_unoccluded'] | |
| image = image_read.copy() | |
| indices = np.random.randint(len(tracks),size=30) | |
| prime_numbers = primes(1000) | |
| unoccluded_index_all = get_unoccluded(indices, tracks) | |
| mask_stacked = image*0 | |
| mask_stacked_all =[] | |
| count = 0 | |
| time = datetime.datetime.now() | |
| for l in indices: | |
| try: | |
| image_crop = get_image(timestamps[l], path) | |
| except: | |
| continue | |
| try: | |
| bb_left, bb_top, bb_width, bb_height, confidence = tracks[l] | |
| except: | |
| bb_left, bb_top, bb_width, bb_height, confidence, track_id = tracks[l] | |
| mask = get_mask(segmentations[l], image) | |
| image[mask > 0] = image_crop[mask > 0] | |
| mask[mask > 0] = 1 | |
| for count, mask_inc in enumerate(mask_stacked_all): | |
| mask_stacked_all[count][cv2.bitwise_and(mask, mask_inc) > 0] = 2 | |
| mask_stacked_all.append(mask) | |
| mask_stacked += mask | |
| count = count+1 | |
| cv2.imwrite(save_path + '/images/'+str(time).replace(' ','T').replace(':','-').split('+')[0] + '.jpg', image[:, :, ::-1]) | |
| cv2.imwrite(save_path + '/Segmentation/'+str(time).replace(' ','T').replace(':','-').split('+')[0] + '.jpg', mask_stacked[:, :, ::-1]*30) | |
| np.savez_compressed(save_path+'/Segmentation/'+str(time).replace(' ','T').replace(':','-').split('+')[0], mask=mask_stacked_all) | |
| def CWALT_Generation(camera_name): | |
| save_path_train = 'data/cwalt_train' | |
| save_path_test = 'data/cwalt_test' | |
| json_file_path = 'data/{}/{}.json'.format(camera_name,camera_name) # iii1/iii1_7_test.json' # './data.json' | |
| path = 'data/' + camera_name | |
| data = np.load(json_file_path + '.npz', allow_pickle=True) | |
| ## slip data | |
| data_train=dict() | |
| data_test=dict() | |
| split_index = int(len(data['timestamps_final_unoccluded'])*0.8) | |
| data_train['tracks_all_unoccluded'] = data['tracks_all_unoccluded'][0:split_index] | |
| data_train['segmentation_all_unoccluded'] = data['segmentation_all_unoccluded'][0:split_index] | |
| data_train['timestamps_final_unoccluded'] = data['timestamps_final_unoccluded'][0:split_index] | |
| data_test['tracks_all_unoccluded'] = data['tracks_all_unoccluded'][split_index:] | |
| data_test['segmentation_all_unoccluded'] = data['segmentation_all_unoccluded'][split_index:] | |
| data_test['timestamps_final_unoccluded'] = data['timestamps_final_unoccluded'][split_index:] | |
| image_read = np.array(Image.open(path + '/T18-median_image.jpg')) | |
| image_read = cv2.resize(image_read, (int(image_read.shape[1]/2), int(image_read.shape[0]/2))) | |
| try: | |
| os.mkdir(save_path_train) | |
| except: | |
| print(save_path_train) | |
| try: | |
| os.mkdir(save_path_train + '/images') | |
| os.mkdir(save_path_train + '/Segmentation') | |
| except: | |
| print(save_path_train+ '/images') | |
| try: | |
| os.mkdir(save_path_test) | |
| except: | |
| print(save_path_test) | |
| try: | |
| os.mkdir(save_path_test + '/images') | |
| os.mkdir(save_path_test + '/Segmentation') | |
| except: | |
| print(save_path_test+ '/images') | |
| for loop in tqdm(range(3000), desc="Generating training CWALT Images "): | |
| save_image(image_read, save_path_train, data_train, path) | |
| for loop in tqdm(range(300), desc="Generating testing CWALT Images "): | |
| save_image(image_read, save_path_test, data_test, path) | |