Showing
4 changed files
with
344 additions
and
0 deletions
source_code/detect.py
0 → 100644
This diff is collapsed. Click to expand it.
source_code/roitrain.py
0 → 100644
| 1 | +from __future__ import division | ||
| 2 | + | ||
| 3 | +from roipool2 import * | ||
| 4 | +from models import * | ||
| 5 | +from utils.utils import * | ||
| 6 | +from utils.datasets import * | ||
| 7 | +from utils.parse_config import * | ||
| 8 | +# from test import evaluate | ||
| 9 | + | ||
| 10 | +from terminaltables import AsciiTable | ||
| 11 | + | ||
| 12 | +import os | ||
| 13 | +import sys | ||
| 14 | +import time | ||
| 15 | +import datetime | ||
| 16 | +import argparse | ||
| 17 | +import warnings | ||
| 18 | + | ||
| 19 | +import torch | ||
| 20 | +from torch.utils.data import DataLoader | ||
| 21 | +from torchvision import datasets | ||
| 22 | +from torchvision import transforms | ||
| 23 | +from torch.autograd import Variable | ||
| 24 | +import torch.optim as optim | ||
| 25 | +import warnings | ||
| 26 | +warnings.filterwarnings("ignore", category=UserWarning) | ||
| 27 | + | ||
| 28 | + | ||
| 29 | + | ||
| 30 | + | ||
| 31 | + | ||
| 32 | +if __name__ == '__main__': | ||
| 33 | + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | ||
| 34 | + print('device: ', device) | ||
| 35 | + | ||
| 36 | + data_config = parse_data_config('config/cafe_distance.data') | ||
| 37 | + train_path = data_config["train"] | ||
| 38 | + valid_path = data_config["valid"] | ||
| 39 | + class_names = load_classes(data_config["names"]) | ||
| 40 | + | ||
| 41 | + model = Darknet('config/yolov3-tiny.cfg', 416).to(device) | ||
| 42 | + | ||
| 43 | + model.load_state_dict(torch.load('checkpoints_cafe_distance/tiny1_2500.pth', map_location=device)) | ||
| 44 | + model.eval() | ||
| 45 | + | ||
| 46 | + dataset = ListDataset(train_path, augment=True, multiscale=True) | ||
| 47 | + dataloader = torch.utils.data.DataLoader( | ||
| 48 | + dataset, | ||
| 49 | + batch_size=1, | ||
| 50 | + shuffle=True, | ||
| 51 | + num_workers=4, | ||
| 52 | + pin_memory=True, | ||
| 53 | + collate_fn=dataset.collate_fn, | ||
| 54 | + ) | ||
| 55 | + | ||
| 56 | + model_distance = ROIPool((3, 3)).to(device) | ||
| 57 | + model_parameters = filter(lambda p: p.requires_grad, model_distance.parameters()) | ||
| 58 | + params = sum([np.prod(p.size()) for p in model_parameters]) | ||
| 59 | + print('Params: ', params) | ||
| 60 | + | ||
| 61 | + optimizer = torch.optim.Adam(model_distance.parameters()) | ||
| 62 | + | ||
| 63 | + | ||
| 64 | + a = [] | ||
| 65 | + for epoch in range(2000): | ||
| 66 | + | ||
| 67 | + warnings.filterwarnings('ignore', category=UserWarning) | ||
| 68 | + for batch_i, (img_path, imgs, targets, targets_distance) in enumerate(dataloader): | ||
| 69 | + | ||
| 70 | + | ||
| 71 | + imgs = Variable(imgs.to(device)) | ||
| 72 | + with torch.no_grad(): | ||
| 73 | + | ||
| 74 | + featuremap, detections = model(imgs) | ||
| 75 | + # print(featuremap.shape) | ||
| 76 | + featuremap = Variable(featuremap.to(device)) | ||
| 77 | + | ||
| 78 | + detections = non_max_suppression(detections, 0.8, 0.4) | ||
| 79 | + targets_distance = torch.tensor(targets_distance[0]) | ||
| 80 | + targets_distance = Variable(targets_distance, requires_grad=True) | ||
| 81 | + | ||
| 82 | + | ||
| 83 | + | ||
| 84 | + if detections is not None: | ||
| 85 | + detections[0] = Variable(detections[0], requires_grad=True) | ||
| 86 | + | ||
| 87 | + | ||
| 88 | + loss, outputs = model_distance(featuremap, detections[0], targets=targets_distance) | ||
| 89 | + # loss = torch.tensor([loss]).to(device) | ||
| 90 | + # loss.requires_grad = True | ||
| 91 | + # print(model_distance.fc1.bias) | ||
| 92 | + optimizer.zero_grad() | ||
| 93 | + loss.backward() | ||
| 94 | + optimizer.step() | ||
| 95 | + # print(model_distance.fc1.bias) | ||
| 96 | + | ||
| 97 | + # print(batch_i) | ||
| 98 | + print(epoch) | ||
| 99 | + | ||
| 100 | + # print(featuremap) | ||
| 101 | + if epoch % 10 == 0: | ||
| 102 | + optimizer.param_groups[0]['lr'] /= 2 | ||
| 103 | + | ||
| 104 | + if epoch % 10 == 0: | ||
| 105 | + torch.save(model_distance.state_dict(), f'checkpoints_distance11/tiny1_{epoch}.pth') |
source_code/train.py
0 → 100644
| 1 | +from __future__ import division | ||
| 2 | + | ||
| 3 | +from models import * | ||
| 4 | +from roipool import * | ||
| 5 | +# from utils.logger import * | ||
| 6 | +from utils.utils import * | ||
| 7 | +from utils.datasets import * | ||
| 8 | +from utils.parse_config import * | ||
| 9 | +# from test import evaluate | ||
| 10 | + | ||
| 11 | +from terminaltables import AsciiTable | ||
| 12 | + | ||
| 13 | +import os | ||
| 14 | +import sys | ||
| 15 | +import time | ||
| 16 | +import datetime | ||
| 17 | +import argparse | ||
| 18 | +import warnings | ||
| 19 | + | ||
| 20 | +import torch | ||
| 21 | +from torch.utils.data import DataLoader | ||
| 22 | +from torchvision import datasets | ||
| 23 | +from torchvision import transforms | ||
| 24 | +from torch.autograd import Variable | ||
| 25 | +import torch.optim as optim | ||
| 26 | +import warnings | ||
| 27 | +warnings.filterwarnings("ignore", category=UserWarning) | ||
| 28 | + | ||
| 29 | +if __name__ == "__main__": | ||
| 30 | + warnings.filterwarnings("ignore", category=UserWarning) | ||
| 31 | + parser = argparse.ArgumentParser() | ||
| 32 | + parser.add_argument("--epochs", type=int, default=8001, help="number of epochs") | ||
| 33 | + parser.add_argument("--batch_size", type=int, default=1, help="size of each image batch") | ||
| 34 | + parser.add_argument("--gradient_accumulations", type=int, default=2, help="number of gradient accums before step") | ||
| 35 | + parser.add_argument("--model_def", type=str, default="config/yolov3-tiny.cfg", help="path to model definition file") | ||
| 36 | + parser.add_argument("--data_config", type=str, default="config/testdata.data", help="path to data config file") | ||
| 37 | + parser.add_argument("--pretrained_weights", type=str, help="if specified starts from checkpoint model") | ||
| 38 | + parser.add_argument("--n_cpu", type=int, default=4, help="number of cpu threads to use during batch generation") | ||
| 39 | + parser.add_argument("--img_size", type=int, default=416, help="size of each image dimension") | ||
| 40 | + parser.add_argument("--checkpoint_interval", type=int, default=50, help="interval between saving model weights") | ||
| 41 | + parser.add_argument("--evaluation_interval", type=int, default=10000, help="interval evaluations on validation set") | ||
| 42 | + parser.add_argument("--compute_map", default=False, help="if True computes mAP every tenth batch") | ||
| 43 | + parser.add_argument("--multiscale_training", default=True, help="allow for multi-scale training") | ||
| 44 | + opt = parser.parse_args() | ||
| 45 | + print(opt) | ||
| 46 | + | ||
| 47 | + # logger = Logger("logs") | ||
| 48 | + | ||
| 49 | + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | ||
| 50 | + print('device: ', device) | ||
| 51 | + | ||
| 52 | + os.makedirs("output", exist_ok=True) | ||
| 53 | + os.makedirs("checkpoints", exist_ok=True) | ||
| 54 | + | ||
| 55 | + # Get data configuration | ||
| 56 | + data_config = parse_data_config(opt.data_config) | ||
| 57 | + train_path = data_config["train"] | ||
| 58 | + valid_path = data_config["valid"] | ||
| 59 | + class_names = load_classes(data_config["names"]) | ||
| 60 | + | ||
| 61 | + # Initiate model | ||
| 62 | + model = Darknet(opt.model_def).to(device) | ||
| 63 | + model.apply(weights_init_normal) | ||
| 64 | + | ||
| 65 | + model_distance = ROIPool((7, 7)).to(device) | ||
| 66 | + | ||
| 67 | + # If specified we start from checkpoint | ||
| 68 | + if opt.pretrained_weights: | ||
| 69 | + if opt.pretrained_weights.endswith(".pth"): | ||
| 70 | + model.load_state_dict(torch.load(opt.pretrained_weights)) | ||
| 71 | + else: | ||
| 72 | + model.load_darknet_weights(opt.pretrained_weights) | ||
| 73 | + | ||
| 74 | + model_parameters = filter(lambda p: p.requires_grad, model.parameters()) | ||
| 75 | + params = sum([np.prod(p.size()) for p in model_parameters]) | ||
| 76 | + print('Params: ', params) | ||
| 77 | + # Get dataloader | ||
| 78 | + dataset = ListDataset(train_path, augment=True, multiscale=opt.multiscale_training) | ||
| 79 | + dataloader = torch.utils.data.DataLoader( | ||
| 80 | + dataset, | ||
| 81 | + batch_size=opt.batch_size, | ||
| 82 | + shuffle=False, | ||
| 83 | + num_workers=opt.n_cpu, | ||
| 84 | + pin_memory=True, | ||
| 85 | + collate_fn=dataset.collate_fn, | ||
| 86 | + ) | ||
| 87 | + | ||
| 88 | + optimizer = torch.optim.Adam(model.parameters()) | ||
| 89 | + | ||
| 90 | + metrics = [ | ||
| 91 | + "grid_size", | ||
| 92 | + "loss", | ||
| 93 | + "x", | ||
| 94 | + "y", | ||
| 95 | + "w", | ||
| 96 | + "h", | ||
| 97 | + "conf", | ||
| 98 | + "cls", | ||
| 99 | + "cls_acc", | ||
| 100 | + "recall50", | ||
| 101 | + "recall75", | ||
| 102 | + "precision", | ||
| 103 | + "conf_obj", | ||
| 104 | + "conf_noobj", | ||
| 105 | + ] | ||
| 106 | + | ||
| 107 | + for epoch in range(opt.epochs): | ||
| 108 | + model.train() | ||
| 109 | + warnings.filterwarnings('ignore', category=UserWarning) | ||
| 110 | + start_time = time.time() | ||
| 111 | + for batch_i, (_, imgs, targets) in enumerate(dataloader): | ||
| 112 | + batches_done = len(dataloader) * epoch + batch_i | ||
| 113 | + | ||
| 114 | + imgs = Variable(imgs.to(device)) | ||
| 115 | + targets = Variable(targets.to(device), requires_grad=False) | ||
| 116 | + | ||
| 117 | + loss, outputs = model(imgs, targets) | ||
| 118 | + print(f'targets = {targets}') | ||
| 119 | + loss.backward() | ||
| 120 | + | ||
| 121 | + if batches_done % opt.gradient_accumulations: | ||
| 122 | + # Accumulates gradient before each step | ||
| 123 | + optimizer.step() | ||
| 124 | + optimizer.zero_grad() | ||
| 125 | + | ||
| 126 | + # ---------------- | ||
| 127 | + # Log progress | ||
| 128 | + # ---------------- | ||
| 129 | + | ||
| 130 | + log_str = "\n---- [Epoch %d/%d, Batch %d/%d] ----\n" % (epoch, opt.epochs, batch_i, len(dataloader)) | ||
| 131 | + | ||
| 132 | + metric_table = [["Metrics", *[f"YOLO Layer {i}" for i in range(len(model.yolo_layers))]]] | ||
| 133 | + | ||
| 134 | + # Log metrics at each YOLO layer | ||
| 135 | + for i, metric in enumerate(metrics): | ||
| 136 | + formats = {m: "%.6f" for m in metrics} | ||
| 137 | + formats["grid_size"] = "%2d" | ||
| 138 | + formats["cls_acc"] = "%.2f%%" | ||
| 139 | + row_metrics = [formats[metric] % yolo.metrics.get(metric, 0) for yolo in model.yolo_layers] | ||
| 140 | + metric_table += [[metric, *row_metrics]] | ||
| 141 | + | ||
| 142 | + # Tensorboard logging | ||
| 143 | + tensorboard_log = [] | ||
| 144 | + for j, yolo in enumerate(model.yolo_layers): | ||
| 145 | + for name, metric in yolo.metrics.items(): | ||
| 146 | + if name != "grid_size": | ||
| 147 | + tensorboard_log += [(f"{name}_{j+1}", metric)] | ||
| 148 | + tensorboard_log += [("loss", loss.item())] | ||
| 149 | + # logger.list_of_scalars_summary(tensorboard_log, batches_done) | ||
| 150 | + | ||
| 151 | + log_str += AsciiTable(metric_table).table | ||
| 152 | + log_str += f"\nTotal loss {loss.item()}" | ||
| 153 | + | ||
| 154 | + # Determine approximate time left for epoch | ||
| 155 | + epoch_batches_left = len(dataloader) - (batch_i + 1) | ||
| 156 | + time_left = datetime.timedelta(seconds=epoch_batches_left * (time.time() - start_time) / (batch_i + 1)) | ||
| 157 | + log_str += f"\n---- ETA {time_left}" | ||
| 158 | + | ||
| 159 | + print(log_str) | ||
| 160 | + | ||
| 161 | + model.seen += imgs.size(0) | ||
| 162 | + | ||
| 163 | + if epoch % opt.evaluation_interval == 0 and epoch != 0: | ||
| 164 | + print("\n---- Evaluating Model ----") | ||
| 165 | + # Evaluate the model on the validation set | ||
| 166 | + precision, recall, AP, f1, ap_class = evaluate( | ||
| 167 | + model, | ||
| 168 | + path=valid_path, | ||
| 169 | + iou_thres=0.5, | ||
| 170 | + conf_thres=0.5, | ||
| 171 | + nms_thres=0.5, | ||
| 172 | + img_size=opt.img_size, | ||
| 173 | + batch_size=1, | ||
| 174 | + ) | ||
| 175 | + evaluation_metrics = [ | ||
| 176 | + ("val_precision", precision.mean()), | ||
| 177 | + ("val_recall", recall.mean()), | ||
| 178 | + ("val_mAP", AP.mean()), | ||
| 179 | + ("val_f1", f1.mean()), | ||
| 180 | + ] | ||
| 181 | + # logger.list_of_scalars_summary(evaluation_metrics, epoch) | ||
| 182 | + | ||
| 183 | + # Print class APs and mAP | ||
| 184 | + ap_table = [["Index", "Class name", "AP"]] | ||
| 185 | + for i, c in enumerate(ap_class): | ||
| 186 | + ap_table += [[c, class_names[c], "%.5f" % AP[i]]] | ||
| 187 | + print(AsciiTable(ap_table).table) | ||
| 188 | + print(f"---- mAP {AP.mean()}") | ||
| 189 | + | ||
| 190 | + if epoch % opt.checkpoint_interval == 0: | ||
| 191 | + torch.save(model.state_dict(), f"checkpoints_fire/tiny1_%d.pth" % (epoch)) |
source_code/video_capture.py
0 → 100644
| 1 | +import cv2 | ||
| 2 | +import queue | ||
| 3 | +import threading | ||
| 4 | + | ||
| 5 | +class BufferlessVideoCapture: | ||
| 6 | + ''' | ||
| 7 | + BufferlessVideoCapture is a wrapper for cv2.VideoCapture, | ||
| 8 | + which doesn't have frame buffer. | ||
| 9 | + @param name: videocapture name | ||
| 10 | + ''' | ||
| 11 | + def __init__(self, name): | ||
| 12 | + self.cap = cv2.VideoCapture(name) | ||
| 13 | + self.q = queue.Queue() | ||
| 14 | + self.thr = threading.Thread(target=self._reader) | ||
| 15 | + self.thr.daemon = True | ||
| 16 | + self.thr.start() | ||
| 17 | + | ||
| 18 | + def _reader(self): | ||
| 19 | + ''' | ||
| 20 | + Main loop for thread. | ||
| 21 | + ''' | ||
| 22 | + while True: | ||
| 23 | + ret, frame = self.cap.read() | ||
| 24 | + if not ret: | ||
| 25 | + break | ||
| 26 | + if not self.q.empty(): | ||
| 27 | + try: | ||
| 28 | + self.q.get_nowait() # discard previous (unprocessed) frame | ||
| 29 | + except queue.Empty: | ||
| 30 | + pass | ||
| 31 | + if self.q.qsize() > 2: | ||
| 32 | + print(self.q.qsize()) | ||
| 33 | + self.q.put(frame) | ||
| 34 | + | ||
| 35 | + def isOpened(self): | ||
| 36 | + return self.cap.isOpened() | ||
| 37 | + | ||
| 38 | + def release(self): | ||
| 39 | + self.cap.release() | ||
| 40 | + | ||
| 41 | + def read(self): | ||
| 42 | + ''' | ||
| 43 | + Read current frame. | ||
| 44 | + ''' | ||
| 45 | + return True, self.q.get() | ||
| 46 | + | ||
| 47 | + def close(self): | ||
| 48 | + pass | ||
| ... | \ No newline at end of file | ... | \ No newline at end of file |
-
Please register or login to post a comment