add preprocess and postprocess

Author: joycenerd

mat2yolo.py

@@ -0,0 +1,83 @@
+import h5py
+import numpy as np
+import argparse
+import os
+from PIL import Image
+
+
+parser=argparse.ArgumentParser()
+parser.add_argument('--data-root',type=str, default='/eva_data/zchin/vrdl_hw2_data',help='data root dir')
+args=parser.parse_args()
+
+
+def get_img_name(f, name_col,idx=0):
+    img_name = ''.join(map(chr, f[name_col[idx][0]][()].flatten()))
+    return(img_name)
+
+
+def get_img_boxes(f, bbox_col,idx=0):
+    """
+    get the 'height', 'left', 'top', 'width', 'label' of bounding boxes of an image
+    :param f: h5py.File
+    :param idx: index of the image
+    :return: dictionary
+    """
+    bbox_prop = ['height', 'left', 'top', 'width', 'label']
+    meta = { key : [] for key in bbox_prop}
+
+    box = f[bbox_col[idx][0]]
+    for key in box.keys():
+        if box[key].shape[0] == 1:
+            meta[key].append(int(box[key][0][0]))
+        else:
+            for i in range(box[key].shape[0]):
+                meta[key].append(int(f[box[key][i][0]][()].item()))
+
+    return meta
+
+def yolo_annot(bbox,img_path,annot_path):
+    annot_f=open(annot_path,'w')
+
+    img=Image.open(img_path)
+    w,h=img.size
+
+    label_cnt=len(bbox['label'])
+    for i in range(label_cnt):
+        label,height,left,top,width=bbox['label'][i],bbox['height'][i],bbox['left'][i],bbox['top'][i],bbox['width'][i]
+        if label==10:
+            label=0
+
+        x_center=left+width/2
+        y_center=top+height/2
+        x_center_norm=x_center/w
+        y_center_norm=y_center/h
+
+        width_norm=width/w
+        height_norm=height/h
+
+        annot_f.write(f'{label} {x_center_norm} {y_center_norm} {width_norm} {height_norm}\n')
+    annot_f.close()
+
+
+
+if __name__=='__main__':
+    annot_dir=os.path.join(args.data_root,'labels/all_train')
+    if not os.path.isdir(annot_dir):
+        os.makedirs(annot_dir)
+        
+    mat_f=os.path.join(args.data_root,'train/digitStruct.mat')
+    mat = h5py.File(mat_f)
+
+    data_size=mat['/digitStruct/name'].shape[0]
+    print(f'Data size: {data_size}')
+    
+    name_col=mat['/digitStruct/name']
+    bbox_col=mat['/digitStruct/bbox']
+    
+    for idx in range(data_size):
+        img_name=get_img_name(mat,name_col,idx)
+        bbox=get_img_boxes(mat,bbox_col,idx)
+        print(img_name, bbox)
+        annot_f=os.path.join(annot_dir,f'{img_name[:-3]}txt')
+        img_path=os.path.join(args.data_root,'train',img_name)
+        yolo_annot(bbox,img_path,annot_f)

train_valid_split.py

@@ -0,0 +1,66 @@
+import glob
+import argparse
+import math
+import random
+import os
+import shutil
+
+
+parser=argparse.ArgumentParser()
+parser.add_argument('--data-root',type=str,default='/eva_data/zchin/vrdl_hw2_data',help='trainig image saving directory')
+parser.add_argument('--ratio',type=float,default=0.2,help='validation data ratio')
+parser.add_argument('--data-dir',type=str,default='./data',help='directory to save train valid split results')
+args=parser.parse_args()
+
+if __name__=='__main__':
+    # if not os.path.isdir(args.data_dir):
+    #     os.makedirs(args.data_dir)
+
+    src_img_dir=os.path.join(args.data_root,'train')
+    data_size = len(glob.glob1(src_img_dir,"*.png"))
+    valid_size=math.floor(data_size*args.ratio)
+
+    img_list=[]
+    for img_path in glob.glob(f'{src_img_dir}/*.png'):
+        img_list.append(img_path)
+
+    idx=random.sample(range(data_size),valid_size)
+
+    # train_rec_path=os.path.join(args.data_dir,'train.txt')
+    # valid_rec_path=os.path.join(args.data_dir,'validation.txt')
+    # train_rec=open(train_rec_path,'w')
+    # valid_rec=open(valid_rec_path,'w')
+    
+    dest_img_dir=os.path.join(args.data_root,'images')
+    train_img_dir=os.path.join(dest_img_dir,'train')
+    valid_img_dir=os.path.join(dest_img_dir,'valid')
+    src_label_dir=os.path.join(args.data_root,'labels/all_train')
+    train_label_dir=src_label_dir.replace('all_train','train')
+    valid_label_dir=src_label_dir.replace('all_train','valid')
+    if not os.path.isdir(dest_img_dir):
+        os.mkdir(dest_img_dir)
+        os.mkdir(train_img_dir)
+        os.mkdir(valid_img_dir)
+        os.mkdir(train_label_dir)
+        os.mkdir(valid_label_dir)
+
+
+    for i in range(data_size):
+        if i in idx:
+            src_img=img_list[i]
+            dest_img=src_img.replace('train','images/valid')
+            shutil.copy(src_img,dest_img)
+            src_label=src_img.replace('train','labels/all_train').replace('png','txt')
+            dest_label=src_label.replace('all_train','valid')
+            shutil.copyfile(src_label,dest_label)
+        else:
+            src_img=img_list[i]
+            dest_img=src_img.replace('train','images/train')
+            shutil.copy(src_img,dest_img)
+            src_label=src_img.replace('train','labels/all_train').replace('png','txt')
+            dest_label=src_label.replace('all_train','train')
+            shutil.copyfile(src_label,dest_label)
+        
+    train_size=len(glob.glob1(train_img_dir,"*.png"))
+    valid_size=len(glob.glob1(valid_img_dir,"*.png"))
+    print(f'train size: {train_size}\tvalid size: {valid_size}')

yolo2coco.py

@@ -0,0 +1,57 @@
+import argparse
+import os
+import cv2
+import json
+
+
+parser=argparse.ArgumentParser()
+parser.add_argument('--yolo-path',type=str,default='yolov5/runs/detect/exp/labels')
+parser.add_argument('--data-root',type=str,default='/eva_data/zchin/vrdl_hw2_data')
+args=parser.parse_args()
+
+
+if __name__=='__main__':
+    img_dir=os.path.join(args.data_root,'test')
+    data_listdir=os.listdir(img_dir)
+    data_listdir.sort(key = lambda x: int(x[:-4]))
+
+    result_to_json=[]
+    for img_name in data_listdir:
+        img_path=os.path.join(img_dir,img_name)
+        img=cv2.imread(img_path)
+        h,w,_=img.shape
+
+        image_id=int(img_name[:-4])
+        label_path=os.path.join(args.yolo_path,img_name.replace('png','txt'))
+        if not os.path.exists(str(label_path)):
+            continue
+        label=open(label_path,'r')
+        lines=label.readlines()
+
+        
+        for line in lines:
+            det_box_info={}
+            
+            det_box_info['image_id']=image_id
+            line=line.strip().split()
+            x_center,y_center,width,height=float(line[1]),float(line[2]),float(line[3]),float(line[4])
+            
+            x_center*=w
+            y_center*=h
+            width*=w
+            height*=h
+            min_x=x_center-width/2
+            min_y=y_center/2
+
+            det_box_info['bbox']=(tuple((min_x,min_y,width,height)))
+            det_box_info['score']=float(line[5])
+            det_box_info['category_id']=int(line[0])
+            result_to_json.append(det_box_info)
+        
+        label.close()
+
+    # Write the list to answer.json 
+    json_object = json.dumps(result_to_json, indent=4)
+
+    with open("answer.json", "w") as outfile:
+        outfile.write(json_object)