├── README.md ├── classes.names ├── extract_frames.py ├── label_img.py ├── params.py └── requirements.txt /README.md: -------------------------------------------------------------------------------- 1 | # Object-Detection-Labeller 2 | A labelling tool which can be used to label objects with the help of a trained model. 3 | 4 | ## Installation 5 | Create a virtual environment and install the requirements. 6 | 7 | ```bash 8 | virtualenv -p /usr/bin/python3.7 venv 9 | 10 | source venv/bin/activate 11 | 12 | pip3 install -r requirements.txt 13 | ``` 14 | 15 | ## Requirements 16 | * Numpy 17 | * OpenCV 18 | * Torch 19 | * Torchvision 20 | * Six 21 | * Pillow 22 | 23 | ## Usage 24 | [Link](https://entuedu-my.sharepoint.com/:u:/g/personal/n1904786f_e_ntu_edu_sg/ESTo9OCxJIlCrWibMQmCrtgBMFDZS0MboObCg_fGPwSzgg?e=nrSZkc) to a faster_rcnn_resnet50_fpn trained on the Bosch Dataset. 25 | 26 | Modify the `img_path`, `save_lbl` and `model_path` in `params.py`. 27 | 28 | `enable_model` flag lets you decide if you want a model to make bounding box predictions which you can choose to keep or discard. 29 | 30 | Only predictions above the set `confidence` will be shown. 31 | 32 | Change `resume` to the name of the image where you left off. 33 | 34 | Run using : 35 | ``` 36 | python3 label_img.py 37 | ``` 38 | 39 | * The model shows every predicted bounding box one by one. 40 | * Press `1-8` based on the `class.names` or `d` to discard the prediction. 41 | * Once all predictions are shown, it switches to manual labelling where the user can add their own bounding boxes and classes. 42 | * Manual mode is indicated by a red `Manual` displayed on top of the screen. 43 | * Click at the top left of the region of interest, drag to the bottom right and release. If you're not satisfied with the bounding box, you can draw again, discarding the previous box. 44 | * After you're satisfied with the box, press `1-8` for the class, which confirms it. 45 | * Repeat for multiple objects. 46 | * Press `s` after you've drawn your bounding boxes to save and go to the next image. 47 | * If the image directly goes manual mode, the model failed to make any predictions. 48 | * Press `q` to quit. 49 | 50 | 51 | ## Save Format 52 | * The co-ordinates are normalized and saved, so the values should be between `0` and `1`. 53 | * `Class_ID X1 Y1 X2 Y2` 54 | -------------------------------------------------------------------------------- /classes.names: -------------------------------------------------------------------------------- 1 | RedLeft 2 | Red 3 | RedRight 4 | GreenLeft 5 | Green 6 | GreenRight 7 | Yellow 8 | off -------------------------------------------------------------------------------- /extract_frames.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | 3 | from params import img_path, vid_path 4 | 5 | cap = cv2.VideoCapture(vid_path) 6 | 7 | f_id = 0 8 | every_n_frames = 1 9 | k = 1 10 | 11 | while True: 12 | ret, img = cap.read() 13 | img_h, img_w, _ = img.shape 14 | f_id += 1 15 | 16 | if f_id % every_n_frames == 0: 17 | cv2.imwrite(img_path + str(k) + '.png', img) 18 | 19 | cv2.imshow("img", cv2.resize(img, (img_w // 2, img_h // 2))) 20 | cv2.waitKey(1) 21 | 22 | print('image' + str(k)) 23 | k += 1 24 | -------------------------------------------------------------------------------- /label_img.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | import cv2 4 | 5 | import torch 6 | import torchvision 7 | from torchvision.transforms import functional as F 8 | 9 | from torchvision.models.detection.faster_rcnn import FastRCNNPredictor 10 | 11 | from params import * 12 | 13 | 14 | def get_coords(event, x, y, flags, param): 15 | global mouseX1, mouseY1, mouseX2, mouseY2, flag 16 | if event == cv2.EVENT_LBUTTONDOWN: 17 | flag = 0 18 | mouseX1, mouseY1 = x, y 19 | 20 | if event == cv2.EVENT_LBUTTONUP: 21 | flag = 1 22 | mouseX2, mouseY2 = x, y 23 | 24 | 25 | if enable_model: 26 | model = torchvision.models.detection.__dict__["fasterrcnn_resnet50_fpn"](pretrained=True) 27 | in_features = model.roi_heads.box_predictor.cls_score.in_features 28 | model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 8) 29 | 30 | device = torch.device("cuda:0") 31 | model.to(device) 32 | 33 | model.eval() 34 | checkpoint = torch.load(model_path, map_location='cpu') 35 | model.load_state_dict(checkpoint['model']) 36 | 37 | 38 | mouseX1, mouseY1 = 0, 0 39 | mouseX2, mouseY2 = 0, 0 40 | 41 | img_list = os.listdir(img_path) 42 | img_list.sort(key=lambda f: int(re.sub('\D', '', f))) 43 | 44 | res_flag = False if resume is "" else True 45 | 46 | cv2.namedWindow('Image') 47 | cv2.setMouseCallback('Image', get_coords) 48 | 49 | for name in img_list: 50 | if resume is not "" and resume in name: 51 | res_flag = False 52 | 53 | if res_flag: 54 | print("Skipping " + name) 55 | continue 56 | 57 | print("Labelling " + name) 58 | save_str = "" 59 | num_boxes = 0 60 | img = cv2.imread(os.path.join(img_path, name)) 61 | img_h, img_w, _ = img.shape 62 | 63 | img = cv2.resize(img, (img_w // 2, img_h // 2)) 64 | img_h, img_w, _ = img.shape 65 | disp = img.copy() 66 | cpy = img.copy() 67 | 68 | if enable_model: 69 | img = F.to_tensor(img).unsqueeze(0) 70 | img = img.to(device) 71 | 72 | pred = model(img) 73 | bboxes = pred[0]['boxes'].cpu().detach().numpy() 74 | scores = pred[0]['scores'].cpu().detach().numpy() 75 | 76 | for bbox, score in zip(bboxes, scores): 77 | save = False 78 | cv2.putText(disp, "Model Prediction", (img_w // 3, 50), 0, 1, (0, 0, 255), 1) 79 | 80 | if score > confidence: 81 | x1, y1, x2, y2 = bbox 82 | 83 | while True: 84 | cv2.rectangle(disp, (x1, y1), (x2, y2), 255, 1) 85 | cv2.imshow("Image", disp) 86 | key = cv2.waitKey(1) 87 | 88 | if key >= ord('0') and key <= ord('9'): 89 | num_boxes += 1 90 | val = key - 48 91 | save = True 92 | break 93 | 94 | if key == ord('d'): 95 | save = False 96 | break 97 | 98 | if key == ord('q'): 99 | print("Stopped at ", name) 100 | exit(0) 101 | 102 | if save: 103 | cv2.rectangle(cpy, (x1, y1), (x2, y2), 255, 1) 104 | disp = cpy.copy() 105 | save_str += str(val - 1) + " " + str(x1 / img_w) + " " + str(y1 / img_h) + " " + str(x2 / img_w) + " " + str(y2 / img_h) + "\n" 106 | 107 | print("Add your own Boxes") 108 | disp = cpy.copy() 109 | cv2.putText(disp, "Manual", (img_w // 2, 50), 0, 1, (0, 0, 255), 1) 110 | cpy = disp.copy() 111 | 112 | while True: 113 | cv2.imshow("Image", disp) 114 | key = cv2.waitKey(1) 115 | 116 | if mouseX1 != 0 and mouseY1 != 0 and mouseX2 != 0 and mouseY2 != 0 and flag: 117 | disp = cpy.copy() 118 | cv2.rectangle(disp, (mouseX1, mouseY1), (mouseX2, mouseY2), (0, 0, 255), 1) 119 | 120 | if key >= ord('0') and key <= ord('9'): 121 | val = key - 48 122 | 123 | if mouseX1 != 0 and mouseY1 != 0 and mouseX2 != 0 and mouseY2 != 0: 124 | num_boxes += 1 125 | save_str += str(val - 1) + " " + str(mouseX1 / img_w) + " " + str(mouseY1 / img_h) + " " + str(mouseX2 / img_w) + " " + str(mouseY2 / img_h) + "\n" 126 | cv2.rectangle(disp, (mouseX1, mouseY1), (mouseX2, mouseY2), (0, 0, 255), 1) 127 | cv2.rectangle(cpy, (mouseX1, mouseY1), (mouseX2, mouseY2), (0, 0, 255), 1) 128 | mouseX1, mouseY1, mouseX1, mouseY2 = 0, 0, 0, 0 129 | 130 | if key == ord('s'): 131 | break 132 | 133 | if key == ord('q'): 134 | print("Stopped at ", name) 135 | exit(0) 136 | 137 | img_id = name[5:-4] 138 | 139 | with open(os.path.join(save_lbl, 'label' + img_id + '.txt'), 'w') as f: 140 | f.writelines(save_str) 141 | 142 | print("Boxes : ", num_boxes) 143 | -------------------------------------------------------------------------------- /params.py: -------------------------------------------------------------------------------- 1 | img_path = "" 2 | save_lbl = "" 3 | model_path = "" 4 | 5 | resume = "" 6 | 7 | enable_model = 1 8 | confidence = 0.6 9 | 10 | vid_path = "" 11 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy==1.18.2 2 | opencv-python==4.2.0.32 3 | Pillow==7.0.0 4 | six==1.14.0 5 | torch==1.4.0 6 | torchvision==0.5.0 7 | --------------------------------------------------------------------------------