|
|
线条不够圆润是因为我为了照顾内存小的电脑,如果要高圆滑的话,就修改下我源码
第一个量化模型 — 你用的是 model_quantized.onnx,量化会损失精度。model.onnx 更准确
第二,修改下采样,先放大 logits 再 argmax,而不是先 argmax 再放大
第三加入点高斯模糊
修改OcclusionDetector.py文件
第一个位置
- self.model_path = model_path or str(
- Path(__file__).parent / 'models' / 'model.onnx')
复制代码 这里模型修改下
第二个位置修改静态方法_logits_to_segmap
- @staticmethod
- def _logits_to_segmap(logits, h, w):
- seg = logits[0].astype(np.float32)
- seg_upscaled = np.zeros((seg.shape[0], h, w), dtype=np.float32)
- for c in range(seg.shape[0]):
- seg_upscaled[c] = cv2.resize(seg[c], (w, h), interpolation=cv2.INTER_LINEAR)
- seg_map = np.argmax(seg_upscaled, axis=0).astype(np.uint8)
- return seg_map
复制代码
第三个地方静态方法mask_to_polys
- def mask_to_polys(mask, min_area=100, epsilon_factor=0.005):
- blurred = cv2.GaussianBlur(mask, (7, 7), 0)
- _, smoothed = cv2.threshold(blurred, 128, 255, cv2.THRESH_BINARY)
- contours, _ = cv2.findContours(smoothed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_L1)
- polys = []
- for contour in contours:
- if cv2.contourArea(contour) < min_area:
- continue
- epsilon = epsilon_factor * cv2.arcLength(contour, True)
- approx = cv2.approxPolyDP(contour, epsilon, True)
- pts = approx.reshape(-1, 2).astype(np.float32)
- if len(pts) >= 3:
- polys.append(pts)
- return polys
复制代码 经过上面三处修改就会变成圆润的了,完整的文件代码是下面的
- import multiprocessing
- import numpy as np
- import cv2
- from pathlib import Path
- class OcclusionDetector:
- CMD_LOAD = 1
- CMD_DETECT = 2
- CMD_SHUTDOWN = 3
- RSP_LOADED = 10
- RSP_RESULT = 11
- RSP_ERROR = 12
- INCLUDE_LABELS = {1, 2, 4, 5, 6, 7, 10, 11, 12}
- EXCLUDE_LABELS = {3, 8, 9, 13, 14, 15, 16, 17, 18}
- HEAD_LABELS = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18}
- MODEL_INPUT_SIZE = 512
- def __init__(self, model_path=None):
- self.model_path = model_path or str(
- Path(__file__).parent / 'models' / 'model.onnx')
- self._req_q = multiprocessing.Queue()
- self._rsp_q = multiprocessing.Queue()
- self._proc = None
- def start(self):
- if self._proc is not None and self._proc.is_alive():
- return
- self._proc = multiprocessing.Process(
- target=self._subprocess_run,
- args=(self._req_q, self._rsp_q, self.model_path),
- daemon=True,
- )
- self._proc.start()
- rsp = self._rsp_q.get(timeout=30)
- if rsp[0] == self.RSP_ERROR:
- raise RuntimeError(f"OcclusionDetector load failed: {rsp[1]}")
- def detect(self, img, landmarks=None):
- if self._proc is None or not self._proc.is_alive():
- self.start()
- self._req_q.put((self.CMD_DETECT, img))
- rsp = self._rsp_q.get(timeout=60)
- if rsp[0] == self.RSP_ERROR:
- raise RuntimeError(f"OcclusionDetector detect failed: {rsp[1]}")
- return rsp[1]
- def shutdown(self):
- if self._proc is not None and self._proc.is_alive():
- self._req_q.put((self.CMD_SHUTDOWN,))
- self._proc.join(timeout=5)
- if self._proc.is_alive():
- self._proc.terminate()
- self._proc = None
- @staticmethod
- def _subprocess_run(req_q, rsp_q, model_path):
- try:
- import onnxruntime as ort
- sess = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])
- rsp_q.put((OcclusionDetector.RSP_LOADED,))
- while True:
- cmd = req_q.get()
- if cmd[0] == OcclusionDetector.CMD_SHUTDOWN:
- break
- elif cmd[0] == OcclusionDetector.CMD_DETECT:
- try:
- img = cmd[1]
- result = OcclusionDetector._run_detection(sess, img)
- rsp_q.put((OcclusionDetector.RSP_RESULT, result))
- except Exception as e:
- rsp_q.put((OcclusionDetector.RSP_ERROR, str(e)))
- except Exception as e:
- rsp_q.put((OcclusionDetector.RSP_ERROR, str(e)))
- @staticmethod
- def _run_detection(sess, img):
- h, w = img.shape[:2]
- input_img = OcclusionDetector._preprocess(img)
- logits = sess.run(None, {'pixel_values': input_img})[0]
- seg_map = OcclusionDetector._logits_to_segmap(logits, h, w)
- return OcclusionDetector._seg_map_to_polys(seg_map, h, w)
- @staticmethod
- def _preprocess(img):
- blob = cv2.dnn.blobFromImage(img, 1.0 / 255.0,
- (OcclusionDetector.MODEL_INPUT_SIZE,
- OcclusionDetector.MODEL_INPUT_SIZE),
- swapRB=True, crop=False)
- return blob.astype(np.float32)
- @staticmethod
- def _logits_to_segmap(logits, h, w):
- seg = logits[0].astype(np.float32)
- seg_upscaled = np.zeros((seg.shape[0], h, w), dtype=np.float32)
- for c in range(seg.shape[0]):
- seg_upscaled[c] = cv2.resize(seg[c], (w, h), interpolation=cv2.INTER_LINEAR)
- seg_map = np.argmax(seg_upscaled, axis=0).astype(np.uint8)
- return seg_map
- @staticmethod
- def _seg_map_to_polys(seg_map, h, w):
- head_mask = np.zeros((h, w), dtype=np.uint8)
- for label in OcclusionDetector.HEAD_LABELS:
- head_mask[seg_map == label] = 255
- include_mask = np.zeros((h, w), dtype=np.uint8)
- for label in OcclusionDetector.INCLUDE_LABELS:
- include_mask[seg_map == label] = 255
- kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
- include_mask = cv2.morphologyEx(include_mask, cv2.MORPH_CLOSE, kernel, iterations=1)
- include_mask = cv2.morphologyEx(include_mask, cv2.MORPH_OPEN, kernel, iterations=1)
- exclude_polys = []
- for label in OcclusionDetector.EXCLUDE_LABELS:
- label_mask = (seg_map == label).astype(np.uint8) * 255
- overlap = cv2.bitwise_and(label_mask, head_mask)
- if np.sum(overlap > 0) == 0:
- continue
- overlap = cv2.morphologyEx(overlap, cv2.MORPH_CLOSE, kernel, iterations=1)
- overlap = cv2.morphologyEx(overlap, cv2.MORPH_OPEN, kernel, iterations=1)
- exclude_polys.extend(
- OcclusionDetector.mask_to_polys(overlap, min_area=80, epsilon_factor=0.003))
- return {
- 'include': OcclusionDetector.mask_to_polys(include_mask, min_area=300, epsilon_factor=0.002),
- 'exclude': exclude_polys,
- }
- @staticmethod
- def mask_to_polys(mask, min_area=100, epsilon_factor=0.005):
- blurred = cv2.GaussianBlur(mask, (7, 7), 0)
- _, smoothed = cv2.threshold(blurred, 128, 255, cv2.THRESH_BINARY)
- contours, _ = cv2.findContours(smoothed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_L1)
- polys = []
- for contour in contours:
- if cv2.contourArea(contour) < min_area:
- continue
- epsilon = epsilon_factor * cv2.arcLength(contour, True)
- approx = cv2.approxPolyDP(contour, epsilon, True)
- pts = approx.reshape(-1, 2).astype(np.float32)
- if len(pts) >= 3:
- polys.append(pts)
- return polys
复制代码
|
|