Spaces:
Running
Running
File size: 1,287 Bytes
49a0323 |
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 |
import os
import re
from sys import argv, exit
import csv
import numpy as np
def natural_sort(l):
convert = lambda text: int(text) if text.isdigit() else text.lower()
alphanum_key = lambda key: [ convert(c) for c in re.split('([0-9]+)', key) ]
return sorted(l, key = alphanum_key)
def getPairs(imgs):
queryIdxs = np.linspace(start=0, stop=len(imgs)-1, num=10).astype(int).tolist()
databaseIdxs = np.linspace(start=10, stop=len(imgs)-10, num=100).astype(int).tolist()
queryImgs = [imgs[idx] for idx in queryIdxs]
databaseImgs = [imgs[idx] for idx in databaseIdxs]
return queryImgs, databaseImgs
def writeCSV(qImgs, dImgs):
with open('rtImagesDepth.csv', 'w', newline='') as file:
writer = csv.writer(file)
title = []
title.append('query')
for i in range(len(dImgs)):
title.append('data' + str(i+1))
writer.writerow(title)
for qImg in qImgs:
row = []
row.append(qImg)
for dImg in dImgs:
row.append(dImg)
writer.writerow(row)
if __name__ == '__main__':
rgbDir = argv[1]
rgbImgs = natural_sort([file for file in os.listdir(rgbDir) if (file.find("jpg") != -1 or file.find("png") != -1)])
rgbImgs = [os.path.join(rgbDir, img) for img in rgbImgs]
queryImgs, databaseImgs = getPairs(rgbImgs)
writeCSV(queryImgs, databaseImgs) |