123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180 |
- import numpy as np
- import cv2
- import random
- from typing import List
- def hex2bgr(hex):
- gmask = 254 << 8
- rmask = 254
- b = hex >> 16
- g = (hex & gmask) >> 8
- r = hex & rmask
- return np.stack([b, g, r]).transpose()
- def union_area(bboxa, bboxb):
- x1 = max(bboxa[0], bboxb[0])
- y1 = max(bboxa[1], bboxb[1])
- x2 = min(bboxa[2], bboxb[2])
- y2 = min(bboxa[3], bboxb[3])
- if y2 < y1 or x2 < x1:
- return -1
- return (y2 - y1) * (x2 - x1)
- def get_yololabel_strings(clslist, labellist):
- content = ''
- for cls, xywh in zip(clslist, labellist):
- content += str(int(cls)) + ' ' + ' '.join([str(e) for e in xywh]) + '\n'
- if len(content) != 0:
- content = content[:-1]
- return content
- # 4 points bbox to 8 points polygon
- def xywh2xyxypoly(xywh, to_int=True):
- xyxypoly = np.tile(xywh[:, [0, 1]], 4)
- xyxypoly[:, [2, 4]] += xywh[:, [2]]
- xyxypoly[:, [5, 7]] += xywh[:, [3]]
- if to_int:
- xyxypoly = xyxypoly.astype(np.int64)
- return xyxypoly
- def xyxy2yolo(xyxy, w: int, h: int):
- if xyxy == [] or xyxy == np.array([]) or len(xyxy) == 0:
- return None
- if isinstance(xyxy, list):
- xyxy = np.array(xyxy)
- if len(xyxy.shape) == 1:
- xyxy = np.array([xyxy])
- yolo = np.copy(xyxy).astype(np.float64)
- yolo[:, [0, 2]] = yolo[:, [0, 2]] / w
- yolo[:, [1, 3]] = yolo[:, [1, 3]] / h
- yolo[:, [2, 3]] -= yolo[:, [0, 1]]
- yolo[:, [0, 1]] += yolo[:, [2, 3]] / 2
- return yolo
- def yolo_xywh2xyxy(xywh: np.array, w: int, h: int, to_int=True):
- if xywh is None:
- return None
- if len(xywh) == 0:
- return None
- if len(xywh.shape) == 1:
- xywh = np.array([xywh])
- xywh[:, [0, 2]] *= w
- xywh[:, [1, 3]] *= h
- xywh[:, [0, 1]] -= xywh[:, [2, 3]] / 2
- xywh[:, [2, 3]] += xywh[:, [0, 1]]
- if to_int:
- xywh = xywh.astype(np.int64)
- return xywh
- def letterbox(im, new_shape=(640, 640), color=(0, 0, 0), auto=False, scaleFill=False, scaleup=True, stride=128):
- # Resize and pad image while meeting stride-multiple constraints
- shape = im.shape[:2] # current shape [height, width]
- if not isinstance(new_shape, tuple):
- new_shape = (new_shape, new_shape)
- # Scale ratio (new / old)
- r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
- if not scaleup: # only scale down, do not scale up (for better val mAP)
- r = min(r, 1.0)
- # Compute padding
- ratio = r, r # width, height ratios
- new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
- dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
- if auto: # minimum rectangle
- dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding
- elif scaleFill: # stretch
- dw, dh = 0.0, 0.0
- new_unpad = (new_shape[1], new_shape[0])
- ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios
- # dw /= 2 # divide padding into 2 sides
- # dh /= 2
- dh, dw = int(dh), int(dw)
- if shape[::-1] != new_unpad: # resize
- im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
- top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
- left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
- im = cv2.copyMakeBorder(im, 0, dh, 0, dw, cv2.BORDER_CONSTANT, value=color) # add border
- return im, ratio, (dw, dh)
- def resize_keepasp(im, new_shape=640, scaleup=True, interpolation=cv2.INTER_LINEAR, stride=None):
- shape = im.shape[:2] # current shape [height, width]
- if new_shape is not None:
- if not isinstance(new_shape, tuple):
- new_shape = (new_shape, new_shape)
- else:
- new_shape = shape
- # Scale ratio (new / old)
- r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
- if not scaleup: # only scale down, do not scale up (for better val mAP)
- r = min(r, 1.0)
- new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
- if stride is not None:
- h, w = new_unpad
- if new_shape[0] % stride != 0:
- new_h = (stride - (new_shape[0] % stride)) + h
- else:
- new_h = h
- if w % stride != 0:
- new_w = (stride - (w % stride)) + w
- else:
- new_w = w
- new_unpad = (new_h, new_w)
- if shape[::-1] != new_unpad: # resize
- im = cv2.resize(im, new_unpad, interpolation=interpolation)
- return im
- def enlarge_window(rect, im_w, im_h, ratio=2.5, aspect_ratio=1.0) -> List:
- assert ratio > 1.0
-
- x1, y1, x2, y2 = rect
- w = x2 - x1
- h = y2 - y1
- # https://numpy.org/doc/stable/reference/generated/numpy.roots.html
- coeff = [aspect_ratio, w+h*aspect_ratio, (1-ratio)*w*h]
- roots = np.roots(coeff)
- roots.sort()
- delta = int(round(roots[-1] / 2 ))
- delta_w = int(delta * aspect_ratio)
- delta_w = min(x1, im_w - x2, delta_w)
- delta = min(y1, im_h - y2, delta)
- rect = np.array([x1-delta_w, y1-delta, x2+delta_w, y2+delta], dtype=np.int64)
- return rect.tolist()
- def draw_connected_labels(num_labels, labels, stats, centroids, names="draw_connected_labels", skip_background=True):
- labdraw = np.zeros((labels.shape[0], labels.shape[1], 3), dtype=np.uint8)
- max_ind = 0
- if isinstance(num_labels, int):
- num_labels = range(num_labels)
- # for ind, lab in enumerate((range(num_labels))):
- for lab in num_labels:
- if skip_background and lab == 0:
- continue
- randcolor = (random.randint(0,255), random.randint(0,255), random.randint(0,255))
- labdraw[np.where(labels==lab)] = randcolor
- maxr, minr = 0.5, 0.001
- maxw, maxh = stats[max_ind][2] * maxr, stats[max_ind][3] * maxr
- minarea = labdraw.shape[0] * labdraw.shape[1] * minr
- stat = stats[lab]
- bboxarea = stat[2] * stat[3]
- if stat[2] < maxw and stat[3] < maxh and bboxarea > minarea:
- pix = np.zeros((labels.shape[0], labels.shape[1]), dtype=np.uint8)
- pix[np.where(labels==lab)] = 255
- rect = cv2.minAreaRect(cv2.findNonZero(pix))
- box = np.int0(cv2.boxPoints(rect))
- labdraw = cv2.drawContours(labdraw, [box], 0, randcolor, 2)
- labdraw = cv2.circle(labdraw, (int(centroids[lab][0]),int(centroids[lab][1])), radius=5, color=(random.randint(0,255), random.randint(0,255), random.randint(0,255)), thickness=-1)
- cv2.imshow(names, labdraw)
- return labdraw
|