diff --git a/detect.py b/detect.py index 4d879b2..6f66d6e 100644 --- a/detect.py +++ b/detect.py @@ -9,16 +9,10 @@ targets_path = "utils/targets_c60.mat" parser = argparse.ArgumentParser() +parser.add_argument("-image_folder", type=str, default="./1047.tif", help="path to images") +parser.add_argument("-output_folder", type=str, default="./output", help="path to outputs") # Get data configuration -if platform == "darwin": # macos - parser.add_argument("-image_folder", type=str, default="./1047.tif", help="path to images") - parser.add_argument("-output_folder", type=str, default="./output", help="path to outputs") - cuda = torch.cuda.is_available() -else: # gcp - parser.add_argument("-image_folder", type=str, default="./1047.tif", help="path to images") - parser.add_argument("-output_folder", type=str, default="./output", help="path to outputs") - cuda = False - +cuda = torch.cuda.is_available() if platform == "darwin" else False # python3 detect.py -plot_flag 1 parser.add_argument("-plot_flag", type=bool, default=True) parser.add_argument("-secondary_classifier", type=bool, default=False) @@ -34,9 +28,9 @@ def detect(opt): if opt.plot_flag: - os.system("rm -rf " + opt.output_folder + "_img") - os.makedirs(opt.output_folder + "_img", exist_ok=True) - os.system("rm -rf " + opt.output_folder) + os.system(f"rm -rf {opt.output_folder}_img") + os.makedirs(f"{opt.output_folder}_img", exist_ok=True) + os.system(f"rm -rf {opt.output_folder}") os.makedirs(opt.output_folder, exist_ok=True) device = torch.device("cuda:0" if cuda else "cpu") @@ -138,7 +132,7 @@ def detect(opt): # pred[:, 1] += y1 # preds.append(pred.unsqueeze(0)) - if len(preds) > 0: + if preds: detections = non_max_suppression( torch.cat(preds, 1), opt.conf_thres, opt.nms_thres, mat_priors, img, model2, device ) @@ -151,7 +145,7 @@ def detect(opt): # Bounding-box colors color_list = [[random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)] for _ in range(len(classes))] - if len(img_detections) == 0: + if not img_detections: return # Iterate through images and save plot of detections @@ -175,10 +169,10 @@ def detect(opt): # write results to .txt file results_path = os.path.join(opt.output_folder, path.split("/")[-1]) - if os.path.isfile(results_path + ".txt"): - os.remove(results_path + ".txt") + if os.path.isfile(f"{results_path}.txt"): + os.remove(f"{results_path}.txt") - results_img_path = os.path.join(opt.output_folder + "_img", path.split("/")[-1]) + results_img_path = os.path.join(f"{opt.output_folder}_img", path.split("/")[-1]) with open(results_path.replace(".bmp", ".tif") + ".txt", "a") as file: for i in unique_classes: n = (detections[:, -1].cpu() == i).sum() @@ -212,7 +206,11 @@ def detect(opt): if opt.plot_flag: from scoring import score - score.score(opt.output_folder + "/", "/Users/glennjocher/Downloads/DATA/xview/xView_train.geojson", ".") + score.score( + f"{opt.output_folder}/", + "/Users/glennjocher/Downloads/DATA/xview/xView_train.geojson", + ".", + ) class ConvNetb(nn.Module): diff --git a/models.py b/models.py index a95809d..56d8c75 100755 --- a/models.py +++ b/models.py @@ -42,7 +42,7 @@ def create_modules(module_defs): elif module_def["type"] == "route": layers = [int(x) for x in module_def["layers"].split(",")] - filters = sum([output_filters[layer_i] for layer_i in layers]) + filters = sum(output_filters[layer_i] for layer_i in layers) modules.add_module("route_%d" % i, EmptyLayer()) elif module_def["type"] == "shortcut": @@ -81,7 +81,7 @@ class YOLOLayer(nn.Module): def __init__(self, anchors, nC, img_dim, anchor_idxs): super(YOLOLayer, self).__init__() - anchors = [(a_w, a_h) for a_w, a_h in anchors] # (pixels) + anchors = list(anchors) nA = len(anchors) self.anchors = anchors @@ -161,7 +161,7 @@ def forward(self, p, targets=None, requestPrecision=False, weight=None, epoch=No # Mask outputs to ignore non-existing objects (but keep confidence predictions) nM = mask.sum().float() - nGT = sum([len(x) for x in targets]) + nGT = sum(len(x) for x in targets) if nM > 0: # wC = weight[torch.argmax(tcls, 1)] # weight class # wC /= sum(wC) diff --git a/utils/datasets.py b/utils/datasets.py index 16586b3..8c862df 100755 --- a/utils/datasets.py +++ b/utils/datasets.py @@ -15,7 +15,7 @@ class ImageFolder: # for eval-only def __init__(self, path, batch_size=1, img_size=416): if os.path.isdir(path): - self.files = sorted(glob.glob("%s/*.*" % path)) + self.files = sorted(glob.glob(f"{path}/*.*")) elif os.path.isfile(path): self.files = [path] @@ -23,7 +23,7 @@ def __init__(self, path, batch_size=1, img_size=416): self.nB = math.ceil(self.nF / batch_size) # number of batches self.batch_size = batch_size self.height = img_size - assert self.nF > 0, "No images found in path %s" % path + assert self.nF > 0, f"No images found in path {path}" # RGB normalization values self.rgb_mean = np.array([60.134, 49.697, 40.746], dtype=np.float32).reshape((3, 1, 1)) @@ -57,11 +57,11 @@ def __len__(self): class ListDataset: # for training def __init__(self, path, batch_size=1, img_size=608, targets_path=""): self.path = path - self.files = sorted(glob.glob("%s/*.bmp" % path)) + self.files = sorted(glob.glob(f"{path}/*.bmp")) self.nF = len(self.files) # number of image files self.nB = math.ceil(self.nF / batch_size) # number of batches self.batch_size = batch_size - assert self.nB > 0, "No images found in path %s" % path + assert self.nB > 0, f"No images found in path {path}" self.height = img_size # load targets self.mat = scipy.io.loadmat(targets_path) @@ -109,7 +109,7 @@ def __next__(self): img_all = [] labels_all = [] - for index, files_index in enumerate(range(ia, ib)): + for files_index in range(ia, ib): # img_path = self.files[self.shuffled_vector[files_index]] # BGR img_path = "%s/%g.bmp" % (self.path, self.shuffled_vector[files_index]) # img_path = '/Users/glennjocher/Downloads/DATA/xview/train_images/2294.bmp' @@ -414,48 +414,44 @@ def random_affine( imw = cv2.warpPerspective( img, M, dsize=(height, height), flags=cv2.INTER_LINEAR, borderValue=borderValue ) # BGR order (YUV-equalized BGR means) - # borderValue = [40.746, 49.697, 60.134]) # RGB - - # Return warped points also - if targets is not None: - if len(targets) > 0: - n = targets.shape[0] - points = targets[:, 1:5].copy() - area0 = (points[:, 2] - points[:, 0]) * (points[:, 3] - points[:, 1]) - - # warp points - xy = np.ones((n * 4, 3)) - xy[:, :2] = points[:, [0, 1, 2, 3, 0, 3, 2, 1]].reshape(n * 4, 2) # x1y1, x2y2, x1y2, x2y1 - xy = (xy @ M.T)[:, :2].reshape(n, 8) - - # create new boxes - x = xy[:, [0, 2, 4, 6]] - y = xy[:, [1, 3, 5, 7]] - xy = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T - - # apply angle-based reduction - radians = a * math.pi / 180 - reduction = max(abs(math.sin(radians)), abs(math.cos(radians))) ** 0.5 - x = (xy[:, 2] + xy[:, 0]) / 2 - y = (xy[:, 3] + xy[:, 1]) / 2 - w = (xy[:, 2] - xy[:, 0]) * reduction - h = (xy[:, 3] - xy[:, 1]) * reduction - xy = np.concatenate((x - w / 2, y - h / 2, x + w / 2, y + h / 2)).reshape(4, n).T - - # reject warped points outside of image - np.clip(xy, 0, height, out=xy) - w = xy[:, 2] - xy[:, 0] - h = xy[:, 3] - xy[:, 1] - area = w * h - ar = np.maximum(w / (h + 1e-16), h / (w + 1e-16)) - i = (w > 4) & (h > 4) & (area / area0 > 0.1) & (ar < 10) - - targets = targets[i] - targets[:, 1:5] = xy[i] - - return imw, targets, M - else: + if targets is None: return imw + if len(targets) > 0: + n = targets.shape[0] + points = targets[:, 1:5].copy() + area0 = (points[:, 2] - points[:, 0]) * (points[:, 3] - points[:, 1]) + + # warp points + xy = np.ones((n * 4, 3)) + xy[:, :2] = points[:, [0, 1, 2, 3, 0, 3, 2, 1]].reshape(n * 4, 2) # x1y1, x2y2, x1y2, x2y1 + xy = (xy @ M.T)[:, :2].reshape(n, 8) + + # create new boxes + x = xy[:, [0, 2, 4, 6]] + y = xy[:, [1, 3, 5, 7]] + xy = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T + + # apply angle-based reduction + radians = a * math.pi / 180 + reduction = max(abs(math.sin(radians)), abs(math.cos(radians))) ** 0.5 + x = (xy[:, 2] + xy[:, 0]) / 2 + y = (xy[:, 3] + xy[:, 1]) / 2 + w = (xy[:, 2] - xy[:, 0]) * reduction + h = (xy[:, 3] - xy[:, 1]) * reduction + xy = np.concatenate((x - w / 2, y - h / 2, x + w / 2, y + h / 2)).reshape(4, n).T + + # reject warped points outside of image + np.clip(xy, 0, height, out=xy) + w = xy[:, 2] - xy[:, 0] + h = xy[:, 3] - xy[:, 1] + area = w * h + ar = np.maximum(w / (h + 1e-16), h / (w + 1e-16)) + i = (w > 4) & (h > 4) & (area / area0 > 0.1) & (ar < 10) + + targets = targets[i] + targets[:, 1:5] = xy[i] + + return imw, targets, M def convert_tif2bmp(p="/Users/glennjocher/Downloads/DATA/xview/val_images_bmp"): @@ -463,11 +459,11 @@ def convert_tif2bmp(p="/Users/glennjocher/Downloads/DATA/xview/val_images_bmp"): import cv2 - files = sorted(glob.glob("%s/*.tif" % p)) + files = sorted(glob.glob(f"{p}/*.tif")) for i, f in enumerate(files): print("%g/%g" % (i + 1, len(files))) img = cv2.imread(f) cv2.imwrite(f.replace(".tif", ".bmp"), img) - os.system("rm -rf " + f) + os.system(f"rm -rf {f}") diff --git a/utils/utils.py b/utils/utils.py index 868ddaf..d0eeae0 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -13,8 +13,7 @@ def load_classes(path): """Loads class labels at 'path'.""" fp = open(path, "r") - names = fp.read().split("\n")[:-1] - return names + return fp.read().split("\n")[:-1] def modelinfo(model): @@ -241,7 +240,7 @@ def xview_feedback_weights(indices): def plot_one_box(x, im, color=None, label=None, line_thickness=None): - tl = line_thickness or round(0.003 * max(im.shape[0:2])) # line thickness + tl = line_thickness or round(0.003 * max(im.shape[:2])) color = color or [random.randint(0, 255) for _ in range(3)] c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3])) cv2.rectangle(im, c1, c2, color, thickness=tl) @@ -295,9 +294,7 @@ def compute_ap(recall, precision): # where X axis (recall) changes value i = np.where(mrec[1:] != mrec[:-1])[0] - # and sum (\Delta recall) * prec - ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1]) - return ap + return np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1]) def bbox_iou(box1, box2, x1y1x2y2=True): @@ -385,9 +382,9 @@ def build_targets(pred_boxes, pred_conf, pred_cls, target, anchor_wh, nA, nC, nG a, gj, gi, t = a[i], gj[i], gi[i], t[i] if len(t.shape) == 1: t = t.view(1, 5) + elif iou_anch_best < 0.10: + continue else: - if iou_anch_best < 0.10: - continue i = 0 tc, gx, gy, gw, gh = t[:, 0].long(), t[:, 1] * nG, t[:, 2] * nG, t[:, 3] * nG, t[:, 4] * nG @@ -429,15 +426,17 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4, mat=None, img """ output = [None for _ in range(len(prediction))] + # Filter out confidence scores below threshold + # Get score and class with highest confidence + + # cross-class NMS --------------------------------------------- + thresh = 0.8 + radius = 30 # area to search for cross-class ious + # Gather bbox priors + srl = 3 # sigma rejection level for image_i, pred in enumerate(prediction): - # Filter out confidence scores below threshold - # Get score and class with highest confidence - - # cross-class NMS --------------------------------------------- - thresh = 0.8 a = pred.clone() a = a[np.argsort(-a[:, 4])] # sort best to worst - radius = 30 # area to search for cross-class ious for i in range(len(a)): if i >= len(a) - 1: break @@ -478,8 +477,6 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4, mat=None, img # if class_prob2[i] > class_prob[i]: # class_pred[i] = class_pred2[i] - # Gather bbox priors - srl = 3 # sigma rejection level mu = mat["class_mu"][class_pred].T sigma = mat["class_sigma"][class_pred].T * srl @@ -544,33 +541,33 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4, mat=None, img max_detections if output[image_i] is None else torch.cat((output[image_i], max_detections)) ) - # # NMS2 - # for c in unique_labels: - # # Get the detections with the particular class - # detections_class = detections[detections[:, -1] == c] - # # Sort the detections by maximum objectness confidence - # _, conf_sort_index = torch.sort(detections_class[:, 4], descending=True) - # detections_class = detections_class[conf_sort_index] - # # Perform non-maximum suppression - # max_detections = [] - # - # while detections_class.shape[0]: - # if len(detections_class) == 1: - # break - # - # ious = bbox_iou(detections_class[0:1], detections_class[1:]) - # - # if ious.max() > 0.5: - # max_detections.append(detections_class[0].unsqueeze(0)) - # - # # Remove detections with IoU >= NMS threshold - # detections_class = detections_class[1:][ious < nms_thres] - # - # if len(max_detections) > 0: - # max_detections = torch.cat(max_detections).data - # # Add max detections to outputs - # output[image_i] = max_detections if output[image_i] is None else torch.cat( - # (output[image_i], max_detections)) + # # NMS2 + # for c in unique_labels: + # # Get the detections with the particular class + # detections_class = detections[detections[:, -1] == c] + # # Sort the detections by maximum objectness confidence + # _, conf_sort_index = torch.sort(detections_class[:, 4], descending=True) + # detections_class = detections_class[conf_sort_index] + # # Perform non-maximum suppression + # max_detections = [] + # + # while detections_class.shape[0]: + # if len(detections_class) == 1: + # break + # + # ious = bbox_iou(detections_class[0:1], detections_class[1:]) + # + # if ious.max() > 0.5: + # max_detections.append(detections_class[0].unsqueeze(0)) + # + # # Remove detections with IoU >= NMS threshold + # detections_class = detections_class[1:][ious < nms_thres] + # + # if len(max_detections) > 0: + # max_detections = torch.cat(max_detections).data + # # Add max detections to outputs + # output[image_i] = max_detections if output[image_i] is None else torch.cat( + # (output[image_i], max_detections)) return output @@ -588,10 +585,14 @@ def secondary_class_detection(x, y, w, h, img, model, device): y2 = np.minimum(y + l, img.shape[0]).astype(np.uint16) n = len(x) - images = [] - for i in range(n): - images.append(cv2.resize(img[y1[i] : y2[i], x1[i] : x2[i]], (height, height), interpolation=cv2.INTER_LINEAR)) - + images = [ + cv2.resize( + img[y1[i] : y2[i], x1[i] : x2[i]], + (height, height), + interpolation=cv2.INTER_LINEAR, + ) + for i in range(n) + ] # # plot # images_numpy = images.copy() # import matplotlib.pyplot as plt @@ -612,7 +613,7 @@ def secondary_class_detection(x, y, w, h, img, model, device): with torch.no_grad(): classes = [] - nB = int(n / 500) + 1 # number of batches + nB = n // 500 + 1 for i in range(nB): # print(i,nB) j0 = int(i * 500)