File manager - Edit - /home/kridsana/webapp.cm.in.th/673190902/u67319090017/remove-bg-web/image_processing.py
Back
from __future__ import annotations import os import uuid from dataclasses import dataclass from typing import Literal, Tuple import cv2 import numpy as np Mode = Literal["threshold", "canny"] @dataclass class Params: mode: Mode = "threshold" blur_ksize: int = 5 # odd: 3,5,7,... threshold: int = 140 # 0..255 (used in threshold mode) canny_low: int = 60 # used in canny mode canny_high: int = 160 # used in canny mode feather: int = 0 # 0=off, else odd like 3,5,7 (soft edge) def _ensure_odd(n: int, min_v: int = 1, max_v: int = 99) -> int: n = max(min_v, min(max_v, int(n))) if n % 2 == 0: n += 1 return n def decode_image(file_bytes: bytes) -> np.ndarray: arr = np.frombuffer(file_bytes, dtype=np.uint8) img = cv2.imdecode(arr, cv2.IMREAD_COLOR) if img is None: raise ValueError("Cannot decode image. Please upload a valid image file.") return img def encode_png(img_bgra: np.ndarray) -> bytes: ok, buf = cv2.imencode(".png", img_bgra) if not ok: raise ValueError("Failed to encode PNG.") return buf.tobytes() def preprocess(img_bgr: np.ndarray, blur_ksize: int) -> Tuple[np.ndarray, np.ndarray]: blur_ksize = _ensure_odd(blur_ksize, 1, 51) gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY) if blur_ksize > 1: gray = cv2.GaussianBlur(gray, (blur_ksize, blur_ksize), 0) return gray, img_bgr def _largest_component(binary_mask: np.ndarray) -> np.ndarray: # binary_mask: uint8 0/255 num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(binary_mask, connectivity=8) if num_labels <= 1: return binary_mask # skip background label 0 areas = stats[1:, cv2.CC_STAT_AREA] largest_idx = 1 + int(np.argmax(areas)) out = np.zeros_like(binary_mask) out[labels == largest_idx] = 255 return out def _clean_mask(mask: np.ndarray) -> np.ndarray: # light post-process for real usability kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)) mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2) mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1) return mask def segment_threshold(gray: np.ndarray, threshold: int) -> np.ndarray: threshold = max(0, min(255, int(threshold))) _, mask = cv2.threshold(gray, threshold, 255, cv2.THRESH_BINARY) # Heuristic: if background is mostly white, invert if np.mean(mask) > 127: mask = cv2.bitwise_not(mask) mask = _clean_mask(mask) mask = _largest_component(mask) return mask def segment_canny(gray: np.ndarray, low: int, high: int) -> np.ndarray: low = max(0, min(255, int(low))) high = max(0, min(255, int(high))) if high < low: low, high = high, low edges = cv2.Canny(gray, low, high) # turn edges into filled region: dilate + close + fill contours edges = cv2.dilate(edges, np.ones((3, 3), np.uint8), iterations=1) edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, np.ones((7, 7), np.uint8), iterations=2) contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) mask = np.zeros_like(edges) if contours: cv2.drawContours(mask, contours, -1, 255, thickness=cv2.FILLED) mask = _clean_mask(mask) mask = _largest_component(mask) return mask def apply_feather(mask: np.ndarray, feather: int) -> np.ndarray: feather = int(feather) if feather <= 1: return mask feather = _ensure_odd(feather, 3, 51) # blur mask to create soft alpha edge soft = cv2.GaussianBlur(mask, (feather, feather), 0) return soft def remove_background(img_bgr: np.ndarray, params: Params) -> Tuple[np.ndarray, np.ndarray]: gray, _ = preprocess(img_bgr, params.blur_ksize) if params.mode == "threshold": mask = segment_threshold(gray, params.threshold) elif params.mode == "canny": mask = segment_canny(gray, params.canny_low, params.canny_high) else: raise ValueError("Unknown mode. Use 'threshold' or 'canny'.") alpha = apply_feather(mask, params.feather) # BGRA output b, g, r = cv2.split(img_bgr) out_bgra = cv2.merge([b, g, r, alpha.astype(np.uint8)]) return out_bgra, mask def save_outputs(out_dir: str, out_bgra: np.ndarray, mask: np.ndarray) -> Tuple[str, str]: os.makedirs(out_dir, exist_ok=True) uid = uuid.uuid4().hex result_path = os.path.join(out_dir, f"result_{uid}.png") mask_path = os.path.join(out_dir, f"mask_{uid}.png") cv2.imwrite(result_path, out_bgra) cv2.imwrite(mask_path, mask) return result_path, mask_path
| ver. 1.4 |
Github
|
.
| PHP 7.4.33 | Generation time: 0.39 |
proxy
|
phpinfo
|
Settings