Spaces:
Sleeping
Sleeping
| """ | |
| Inference-time preprocessing สำหรับ Space | |
| - Crop มือด้วย ianpan/bone-age-crop | |
| - Histogram matching กับ ref_img.png (ถ้ามี) | |
| - Resize (LongestMaxSize=512) + Pad ให้เป็น 512x512 | |
| ให้ผลลัพธ์เป็น grayscale uint8 512x512 (เหมือน val/test ตอนเทรน) | |
| """ | |
| import os | |
| import cv2 | |
| import torch | |
| import numpy as np | |
| import albumentations as A | |
| from skimage.exposure import match_histograms | |
| from transformers import AutoModel | |
| _CROP_MODEL = None | |
| _REF_IMG = None | |
| _DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| REF_IMG_PATH = os.path.join(os.path.dirname(__file__), "ref_img.png") | |
| _resize_pad = A.Compose([ | |
| A.LongestMaxSize(max_size=512, p=1), | |
| A.PadIfNeeded(512, 512, border_mode=cv2.BORDER_CONSTANT, value=0, p=1), | |
| ]) | |
| def _get_crop_model(): | |
| global _CROP_MODEL | |
| if _CROP_MODEL is None: | |
| _CROP_MODEL = AutoModel.from_pretrained( | |
| "ianpan/bone-age-crop", trust_remote_code=True | |
| ).eval().to(_DEVICE) | |
| return _CROP_MODEL | |
| def _get_ref_img(): | |
| global _REF_IMG | |
| if _REF_IMG is None and os.path.exists(REF_IMG_PATH): | |
| _REF_IMG = cv2.imread(REF_IMG_PATH, 0) | |
| return _REF_IMG | |
| def preprocess_image(image: np.ndarray) -> np.ndarray: | |
| """ | |
| image: grayscale (H, W) uint8 | |
| return: grayscale (512, 512) uint8 | |
| """ | |
| if image.ndim == 3: | |
| image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) | |
| image = image.astype(np.uint8) | |
| # --- 1) Crop --- | |
| try: | |
| crop_model = _get_crop_model() | |
| img_shape = torch.tensor([image.shape[:2]]).to(_DEVICE) | |
| x_crop = crop_model.preprocess(image) | |
| x_crop = torch.from_numpy(x_crop).unsqueeze(0).unsqueeze(0).float().to(_DEVICE) | |
| with torch.inference_mode(): | |
| coords = crop_model(x_crop, img_shape) | |
| x_c, y_c, w_c, h_c = coords[0].cpu().numpy() | |
| cropped = image[int(y_c): int(y_c + h_c), int(x_c): int(x_c + w_c)] | |
| if cropped.size == 0: | |
| cropped = image | |
| except Exception as e: | |
| import traceback | |
| print("=" * 60) | |
| print("[crop] FAILED -> ใช้ภาพเต็มเฟรมแทน (ไม่ได้ crop!)") | |
| print(f"[crop] error: {repr(e)}") | |
| traceback.print_exc() | |
| print("=" * 60) | |
| cropped = image | |
| # --- 2) Histogram matching (ถ้ามี ref_img) --- | |
| ref = _get_ref_img() | |
| if ref is not None: | |
| matched = match_histograms(cropped, ref) | |
| matched = np.clip(matched, 0, 255).astype(np.uint8) | |
| else: | |
| print("[preprocess] ⚠️ ไม่พบ ref_img.png -> ข้าม histogram matching (ผลอาจเพี้ยน)") | |
| matched = cropped | |
| # --- 3) Resize + Pad --- | |
| out = _resize_pad(image=matched)["image"] | |
| return out | |