dk - edit

2026-07-09 18:16:06 +09:00
parent dda1759322
commit 4f1c320b3b
25 changed files with 1744 additions and 4 deletions
@@ -0,0 +1,96 @@
+import os
+from shutil import copyfile
+
+def pytorch_prepare(data_dir, dataset_name):
+    dataset_dir = os.path.join(data_dir, dataset_name)
+
+    if not os.path.isdir(dataset_dir):
+        print('please change the download_path')
+
+    pytorch_path = os.path.join(dataset_dir , 'pytorch')
+
+    if not os.path.isdir(pytorch_path):
+        os.mkdir(pytorch_path)
+        #-----------------------------------------
+        #query
+        print('generatring ' + dataset_name + ' query images.')
+        query_dir = os.path.join(dataset_dir , 'query')
+        query_save_dir = os.path.join(dataset_dir , 'pytorch', 'query')
+        if not os.path.isdir(query_save_dir):
+            os.mkdir(query_save_dir)
+
+        for root, dirs, files in os.walk(query_dir, topdown=True):
+            for name in files:
+                if not name[-3:]=='jpg':
+                    continue
+                ID  = name.split('_')
+                src_dir = os.path.join(query_dir , name)
+                dst_dir = os.path.join(query_save_dir, ID[0])
+                if not os.path.isdir(dst_dir):
+                    os.mkdir(dst_dir)
+                copyfile(src_dir, os.path.join(dst_dir , name))
+        #-----------------------------------------
+        #gallery
+        print('generatring '+dataset_name+' gallery images.')
+        gallery_dir = os.path.join(dataset_dir , 'bounding_box_test')
+        gallery_save_dir = os.path.join(dataset_dir , 'pytorch' , 'gallery')
+        if not os.path.isdir(gallery_save_dir):
+            os.mkdir(gallery_save_dir)
+
+        for root, dirs, files in os.walk(gallery_dir, topdown=True):
+            for name in files:
+                if not name[-3:]=='jpg':
+                    continue
+                ID  = name.split('_')
+                src_dir = os.path.join(gallery_dir, name)
+                dst_dir = os.path.join(gallery_save_dir, ID[0])
+                if not os.path.isdir(dst_dir):
+                    os.mkdir(dst_dir)
+                copyfile(src_dir, os.path.join(dst_dir,name))
+        #---------------------------------------
+        #train_all
+        print('generatring '+dataset_name + ' all training images.')
+        train_dir = os.path.join( dataset_dir , 'bounding_box_train')
+        train_save_all_dir = os.path.join( dataset_dir , 'pytorch', 'train_all')
+        if not os.path.isdir(train_save_all_dir):
+            os.mkdir(train_save_all_dir)
+
+        for root, dirs, files in os.walk(train_dir, topdown=True):
+            for name in files:
+                if not name[-3:]=='jpg':
+                    continue
+                ID  = name.split('_')
+                src_dir = os.path.join(train_dir , name)
+                dst_dir = os.path.join(train_save_all_dir, ID[0])
+                if not os.path.isdir(dst_dir):
+                    os.mkdir(dst_dir)
+                copyfile(src_dir, os.path.join(dst_dir, name))
+                
+        #---------------------------------------
+        #train_val
+        print('generatring '+ dataset_name+' training and validation images.')
+        train_save_dir = os.path.join(dataset_dir, 'pytorch', 'train')
+        val_save_dir = os.path.join(dataset_dir , 'pytorch' , 'val')
+        if not os.path.isdir(train_save_dir):
+            os.mkdir(train_save_dir)
+            os.mkdir(val_save_dir)
+
+        for root, dirs, files in os.walk(train_dir, topdown=True):
+            for name in files:
+                if not name[-3:]=='jpg':
+                    continue
+                ID  = name.split('_')
+                src_dir = os.path.join(train_dir , name)
+                dst_dir = os.path.join(train_save_dir , ID[0])
+                if not os.path.isdir(dst_dir):
+                    os.mkdir(dst_dir)
+                    dst_dir = os.path.join(val_save_dir, ID[0])  #first image is used as val image
+                    os.mkdir(dst_dir)
+                copyfile(src_dir, os.path.join(dst_dir , name))
+        print('Finished ' + dataset_name)
+    else:
+        print(dataset_name + ' pytorch directory exists!')
+
+def pytorch_prepare_all(data_dir):
+    pytorch_prepare('Market1501', data_dir)
+    pytorch_prepare('DukeMTMC', data_dir)