update the modal test
This commit is contained in:
DatTT127
@@ -1,852 +0,0 @@
|
||||
from fastapi import FastAPI, File, UploadFile, HTTPException, Query, Request, Response
|
||||
from pydantic import BaseModel
|
||||
from fastapi.responses import JSONResponse, FileResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
import os
|
||||
import uuid
|
||||
import numpy as np
|
||||
from PIL import Image, ImageDraw
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torchvision.models as models
|
||||
from torchvision import transforms
|
||||
import uvicorn
|
||||
from pathlib import Path
|
||||
import cv2
|
||||
import io
|
||||
import base64
|
||||
from datetime import datetime
|
||||
import re
|
||||
|
||||
# Import custom models
|
||||
import sys
|
||||
sys.path.append('.')
|
||||
from arch.efficientfeedback import EfficientFeedbackNetwork
|
||||
from arch.unet3plus_att import UNet3Plus_Attention
|
||||
from pdf_service import generate_medical_report
|
||||
|
||||
# Configuration
|
||||
UPLOAD_FOLDER = 'uploads'
|
||||
RESULTS_FOLDER = 'results'
|
||||
TEMPLATES_FOLDER = 'templates'
|
||||
|
||||
# for folder in [UPLOAD_FOLDER, RESULTS_FOLDER, TEMPLATES_FOLDER]:
|
||||
# os.makedirs(folder, exist_ok=True)
|
||||
os.makedirs(TEMPLATES_FOLDER, exist_ok=True)
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
# Classes
|
||||
ANGLE_CLASSES = ['med-lat', 'post-trans', 'sup-trans-flex', 'sup-up-long']
|
||||
SEGMENT_CLASSES_SUPRAPAT = {0: "background", 1: "effusion", 2: "fat", 3: "fat-pat", 4: "femur", 5: "synovium", 6: "tendon"}
|
||||
SEGMENT_CLASSES_POST = {0: "background", 1: "fat", 2: "tendon", 3: "muscle", 4: "femur", 5: "artery", 6: "baker's cyst"}
|
||||
|
||||
# Color map for Suprapat
|
||||
COLOR_MAP_SUP = {
|
||||
'background': {'color': [0, 0, 0], 'name': 'Nền'},
|
||||
'effusion': {'color': [255, 0, 0], 'name': 'Dịch khớp'},
|
||||
'fat': {'color': [255, 255, 0], 'name': 'Mỡ'},
|
||||
'fat-pat': {'color': [0, 255, 255], 'name': 'Mỡ Hoffa'},
|
||||
'femur': {'color': [0, 255, 0], 'name': 'Xương đùi'},
|
||||
'synovium': {'color': [255, 0, 255], 'name': 'Màng hoạt dịch'},
|
||||
'tendon': {'color': [0, 0, 255], 'name': 'Gân'}
|
||||
}
|
||||
|
||||
# Color map for Post-trans
|
||||
COLOR_MAP_POST = {
|
||||
'background': {'color': [0, 0, 0], 'name': 'Nền'},
|
||||
"baker's cyst": {'color': [255, 0, 0], 'name': "Baker's cyst"},
|
||||
'fat': {'color': [255, 255, 0], 'name': 'Mỡ'},
|
||||
'muscle': {'color': [0, 255, 255], 'name': 'Cơ bắp'},
|
||||
'femur': {'color': [0, 255, 0], 'name': 'Xương đùi'},
|
||||
'synovium': {'color': [255, 0, 255], 'name': 'Màng hoạt dịch'},
|
||||
'tendon': {'color': [0, 0, 255], 'name': 'Gân'}
|
||||
}
|
||||
|
||||
# Measurement configuration
|
||||
DEFAULT_MEASURE_IDS = [1, 5]
|
||||
PIXEL_TO_MM = 45.0 / 655.0
|
||||
|
||||
app = FastAPI(title="Medical Image Analysis API")
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# app.mount("/uploads", StaticFiles(directory=UPLOAD_FOLDER), name="uploads")
|
||||
# app.mount("/results", StaticFiles(directory=RESULTS_FOLDER), name="results")
|
||||
|
||||
# ============ MODEL LOADING ============
|
||||
|
||||
def load_angle_model(model_name):
|
||||
print(f"📄 Loading angle model: {model_name}")
|
||||
|
||||
if model_name == "convnext":
|
||||
model = models.convnext_tiny(weights=None)
|
||||
model.classifier[2] = nn.Linear(model.classifier[2].in_features, 4)
|
||||
checkpoint = torch.load(f"models/best_convnext_tiny.pth", map_location=device, weights_only=False)
|
||||
checkpoint = {k.replace('model.', ''): v for k, v in checkpoint.items()}
|
||||
model.load_state_dict(checkpoint)
|
||||
elif model_name == "densenet":
|
||||
model = models.densenet121(weights=None)
|
||||
model.classifier = nn.Linear(model.classifier.in_features, 4)
|
||||
checkpoint = torch.load(f"models/best_densenet.pth", map_location=device, weights_only=False)
|
||||
checkpoint = {k.replace('model.', ''): v for k, v in checkpoint.items()}
|
||||
model.load_state_dict(checkpoint)
|
||||
elif model_name == "resnet50":
|
||||
model = models.resnet50(weights=None)
|
||||
model.fc = nn.Linear(model.fc.in_features, 4)
|
||||
checkpoint = torch.load(f"models/best_resnet50.pth", map_location=device, weights_only=False)
|
||||
checkpoint = {k.replace('model.', ''): v for k, v in checkpoint.items()}
|
||||
model.load_state_dict(checkpoint)
|
||||
elif model_name == "efficientnet_b2":
|
||||
model = models.efficientnet_b2(weights=None)
|
||||
model.classifier[1] = nn.Linear(model.classifier[1].in_features, 4)
|
||||
checkpoint = torch.load(f"models/best_efficientnet_b2.pth", map_location=device, weights_only=False)
|
||||
checkpoint = {k.replace('model.', ''): v for k, v in checkpoint.items()}
|
||||
model.load_state_dict(checkpoint)
|
||||
elif model_name == "swin":
|
||||
model = models.swin_v2_s(weights=None)
|
||||
model.head = nn.Linear(model.head.in_features, 4)
|
||||
checkpoint = torch.load(f"models/best_swin_v2_s.pth", map_location=device, weights_only=False)
|
||||
checkpoint = {k.replace('model.', ''): v for k, v in checkpoint.items()}
|
||||
model.load_state_dict(checkpoint)
|
||||
else:
|
||||
raise ValueError(f"Unknown angle model: {model_name}")
|
||||
|
||||
print(f"✅ Loaded: {model_name}")
|
||||
return model.to(device).eval()
|
||||
|
||||
def load_inflammation_model():
|
||||
print("📄 Loading inflammation model")
|
||||
model = models.efficientnet_b0(weights=None)
|
||||
model.classifier[1] = nn.Linear(model.classifier[1].in_features, 2)
|
||||
model.load_state_dict(torch.load("models/efficientnet_b0_ultrasound_2_class.pth", map_location=device, weights_only=False))
|
||||
print("✅ Loaded inflammation model")
|
||||
return model.to(device).eval()
|
||||
|
||||
def load_segmentation_model_sup(model_name):
|
||||
print(f"📄 Loading segmentation model SUP: {model_name}")
|
||||
|
||||
if model_name == "deeplabv3":
|
||||
model = models.segmentation.deeplabv3_resnet50(weights=None)
|
||||
in_ch = model.classifier[-1].in_channels
|
||||
model.classifier = nn.Sequential(
|
||||
model.classifier[0],
|
||||
nn.Dropout(0.3),
|
||||
nn.Conv2d(in_ch, 7, kernel_size=1)
|
||||
)
|
||||
model.load_state_dict(torch.load("models/best_model_Deeplav3.pth", map_location=device, weights_only=False), strict=False)
|
||||
elif model_name == "unet_resnet101":
|
||||
try:
|
||||
from segmentation_models_pytorch import Unet
|
||||
model = Unet(encoder_name="resnet101", encoder_weights=None, classes=7)
|
||||
model.load_state_dict(torch.load("models/unet_resnet101.pth", map_location=device, weights_only=False))
|
||||
except ImportError:
|
||||
raise ValueError("segmentation_models_pytorch not installed")
|
||||
elif model_name == "efficientfeedback":
|
||||
model = EfficientFeedbackNetwork(in_channels=3, num_class=7)
|
||||
model.load_state_dict(torch.load("models/efficientfeedback.pth", map_location=device, weights_only=False))
|
||||
elif model_name == "unet3plus":
|
||||
model = UNet3Plus_Attention(in_channels=3, num_classes=7)
|
||||
model.load_state_dict(torch.load("models/unet3plus_att.pth", map_location=device, weights_only=False))
|
||||
else:
|
||||
raise ValueError(f"Unknown segmentation model: {model_name}")
|
||||
|
||||
print(f"✅ Loaded SUP: {model_name}")
|
||||
return model.to(device).eval()
|
||||
|
||||
def load_segmentation_model_post(model_name):
|
||||
print(f"📄 Loading segmentation model POST: {model_name}")
|
||||
|
||||
if model_name == "deeplabv3_resnet101":
|
||||
model = models.segmentation.deeplabv3_resnet101(weights=None)
|
||||
in_ch = model.classifier[-1].in_channels
|
||||
model.classifier = nn.Sequential(
|
||||
model.classifier[0],
|
||||
nn.Dropout(0.3),
|
||||
nn.Conv2d(in_ch, 7, kernel_size=1)
|
||||
)
|
||||
model.load_state_dict(torch.load("models/best_model_deeplabv3_resnet101_seed_16.pth", map_location=device, weights_only=False), strict=False)
|
||||
else:
|
||||
raise ValueError(f"Unknown post segmentation model: {model_name}")
|
||||
|
||||
print(f"✅ Loaded POST: {model_name}")
|
||||
return model.to(device).eval()
|
||||
|
||||
# Transforms
|
||||
angle_transform = transforms.Compose([
|
||||
transforms.Resize((224, 224)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
|
||||
])
|
||||
|
||||
inflammation_transform = transforms.Compose([
|
||||
transforms.Resize((224, 224)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
|
||||
])
|
||||
|
||||
segmentation_transform = transforms.Compose([
|
||||
transforms.Resize((512, 512)),
|
||||
transforms.ToTensor()
|
||||
])
|
||||
|
||||
# ============ PREDICTION ============
|
||||
|
||||
@torch.no_grad()
|
||||
def predict_angle(model, image_pil):
|
||||
img_tensor = angle_transform(image_pil).unsqueeze(0).to(device)
|
||||
output = model(img_tensor)
|
||||
probs = torch.softmax(output, dim=1)
|
||||
pred_class = torch.argmax(probs, dim=1).item()
|
||||
confidence = probs[0][pred_class].item()
|
||||
return ANGLE_CLASSES[pred_class], round(confidence * 100, 2)
|
||||
|
||||
@torch.no_grad()
|
||||
def predict_inflammation(model, image_pil):
|
||||
img_tensor = inflammation_transform(image_pil).unsqueeze(0).to(device)
|
||||
output = model(img_tensor)
|
||||
probs = torch.softmax(output, dim=1)
|
||||
pred_class = torch.argmax(probs, dim=1).item()
|
||||
confidence = probs[0][pred_class].item()
|
||||
is_inflammation = (pred_class == 1)
|
||||
return is_inflammation, round(confidence * 100, 2)
|
||||
|
||||
@torch.no_grad()
|
||||
def segment_image(model, image_pil, model_type, angle_type):
|
||||
original_size = image_pil.size
|
||||
img_tensor = segmentation_transform(image_pil).unsqueeze(0).to(device)
|
||||
|
||||
if model_type.startswith("deeplabv3"):
|
||||
outputs = model(img_tensor)['out']
|
||||
else:
|
||||
outputs = model(img_tensor)
|
||||
|
||||
upsampled = F.interpolate(outputs, size=original_size[::-1], mode='bilinear', align_corners=False)
|
||||
preds = upsampled.argmax(dim=1)[0].cpu().numpy()
|
||||
|
||||
if angle_type == 'sup' and model_type in ["unet3plus", "efficientfeedback"]:
|
||||
remap = {0: 0, 1: 2, 2: 6, 3: 1, 4: 4, 5: 5, 6: 3}
|
||||
preds = np.vectorize(remap.get)(preds)
|
||||
|
||||
class_map = SEGMENT_CLASSES_SUPRAPAT if angle_type == 'sup' else SEGMENT_CLASSES_POST
|
||||
|
||||
masks = {}
|
||||
for class_id, class_name in class_map.items():
|
||||
mask = (preds == class_id).astype(np.uint8)
|
||||
masks[class_name] = mask
|
||||
|
||||
return preds, masks
|
||||
|
||||
def get_mask_bounding_box(mask, dist_percent=0.01):
|
||||
"""
|
||||
Duyệt toàn bộ vùng được mask, loại bỏ nhiễu và trả về khung bao (Bounding Box).
|
||||
Áp dụng quy tắc kết hợp:
|
||||
1. Giữ lại khối có diện tích lớn nhất (vùng trung tâm).
|
||||
2. Giữ lại các khối phụ nếu thỏa mãn một trong hai điều kiện:
|
||||
- Diện tích >= 1/5 diện tích khối lớn nhất.
|
||||
- Khoảng cách tới khối lớn nhất <= dist_percent * chiều rộng ảnh.
|
||||
"""
|
||||
if mask is None or np.sum(mask) == 0:
|
||||
return None
|
||||
|
||||
# 1. Chuyển sang uint8
|
||||
mask_uint8 = mask.astype(np.uint8)
|
||||
if np.max(mask_uint8) == 1:
|
||||
mask_uint8 *= 255
|
||||
|
||||
# Lấy chiều rộng ảnh để tính ngưỡng khoảng cách theo %
|
||||
img_width = mask_uint8.shape[1]
|
||||
dist_threshold = img_width * dist_percent
|
||||
|
||||
# 2. Làm sạch mask cơ bản (Morphological Opening)
|
||||
kernel = np.ones((5, 5), np.uint8)
|
||||
clean_mask = cv2.morphologyEx(mask_uint8, cv2.MORPH_OPEN, kernel)
|
||||
|
||||
# 3. Tìm các đường bao (các khối tách rời)
|
||||
contours, _ = cv2.findContours(clean_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
|
||||
if not contours:
|
||||
return None
|
||||
|
||||
# 4. Tính diện tích từng khối và tìm khối lớn nhất
|
||||
contour_info = []
|
||||
for cnt in contours:
|
||||
contour_info.append({'cnt': cnt, 'area': cv2.contourArea(cnt)})
|
||||
|
||||
# Sắp xếp theo diện tích giảm dần
|
||||
contour_info.sort(key=lambda x: x['area'], reverse=True)
|
||||
main_block = contour_info[0]
|
||||
max_area = main_block['area']
|
||||
|
||||
if max_area < 50:
|
||||
return None
|
||||
|
||||
# 5. Chuẩn bị để tính khoảng cách (Distance Transform)
|
||||
main_mask = np.zeros_like(mask_uint8)
|
||||
cv2.drawContours(main_mask, [main_block['cnt']], -1, 255, -1)
|
||||
# dist_map chứa khoảng cách từ mỗi điểm tới biên gần nhất của khối chính
|
||||
dist_map = cv2.distanceTransform(255 - main_mask, cv2.DIST_L2, 3)
|
||||
|
||||
# 6. Lọc các khối
|
||||
significant_contours = [main_block['cnt']]
|
||||
area_threshold = max_area / 4.0
|
||||
|
||||
for i in range(1, len(contour_info)):
|
||||
other = contour_info[i]
|
||||
|
||||
# Tạo mask cho khối đang xét để lấy giá trị khoảng cách
|
||||
other_mask = np.zeros_like(mask_uint8)
|
||||
cv2.drawContours(other_mask, [other['cnt']], -1, 255, -1)
|
||||
|
||||
# Khoảng cách nhỏ nhất từ khối này tới khối chính
|
||||
min_dist = np.min(dist_map[other_mask > 0])
|
||||
|
||||
# Điều kiện giữ lại: (Diện tích đủ lớn) HOẶC (Ở gần khối chính theo %)
|
||||
if other['area'] >= area_threshold or min_dist <= dist_threshold:
|
||||
significant_contours.append(other['cnt'])
|
||||
|
||||
# 7. Tính toán bounding box bao quanh tất cả các vùng được chọn
|
||||
all_points = np.concatenate(significant_contours)
|
||||
x, y, w, h = cv2.boundingRect(all_points)
|
||||
|
||||
return x, y, w, h
|
||||
|
||||
def find_max_continuous_segment(col_array):
|
||||
padded = np.concatenate(([0], col_array, [0]))
|
||||
diffs = np.diff(padded)
|
||||
|
||||
starts = np.where(diffs == 1)[0]
|
||||
ends = np.where(diffs == -1)[0]
|
||||
|
||||
if len(starts) == 0:
|
||||
return 0, -1, -1
|
||||
|
||||
lengths = ends - starts
|
||||
max_idx = np.argmax(lengths)
|
||||
max_len = lengths[max_idx]
|
||||
|
||||
return max_len, starts[max_idx], ends[max_idx]
|
||||
|
||||
def measure_thickness_new(masks, image_size, measure_ids=None):
|
||||
if measure_ids is None:
|
||||
measure_ids = DEFAULT_MEASURE_IDS
|
||||
|
||||
width, height = image_size
|
||||
|
||||
mask_all_labels = np.zeros((height, width), dtype=np.uint8)
|
||||
mask_measure = np.zeros((height, width), dtype=np.uint8)
|
||||
|
||||
has_any_label = False
|
||||
|
||||
if 'fat-pat' in masks:
|
||||
class_map = SEGMENT_CLASSES_SUPRAPAT
|
||||
else:
|
||||
class_map = SEGMENT_CLASSES_POST
|
||||
|
||||
for class_id, class_name in class_map.items():
|
||||
if class_name not in masks or class_name == 'background':
|
||||
continue
|
||||
mask = masks[class_name]
|
||||
if np.sum(mask) > 0:
|
||||
has_any_label = True
|
||||
mask_all_labels = np.logical_or(mask_all_labels, mask).astype(np.uint8)
|
||||
|
||||
if class_id in measure_ids:
|
||||
mask_measure = np.logical_or(mask_measure, mask).astype(np.uint8)
|
||||
|
||||
if not has_any_label or np.sum(mask_measure) == 0:
|
||||
return None
|
||||
|
||||
# Đóng khung toàn bộ vùng được mask (xương, màng, dịch, mỡ...)
|
||||
bbox_all = get_mask_bounding_box(mask_all_labels)
|
||||
if bbox_all is None:
|
||||
return None
|
||||
|
||||
x_all, y_all, w_all, h_all = bbox_all
|
||||
|
||||
# Từ khung này, xác định vùng quét là 1/3 ở giữa theo chiều ngang
|
||||
roi_start = x_all + (w_all // 3)
|
||||
roi_end = x_all + (2 * w_all // 3)
|
||||
|
||||
roi_strip = mask_measure[:, roi_start:roi_end]
|
||||
|
||||
global_max_len_px = 0
|
||||
best_x_rel = 0
|
||||
best_y_start = 0
|
||||
best_y_end = 0
|
||||
|
||||
for x in range(roi_strip.shape[1]):
|
||||
col = roi_strip[:, x]
|
||||
if not np.any(col):
|
||||
continue
|
||||
|
||||
length, y_s, y_e = find_max_continuous_segment(col)
|
||||
|
||||
if length > global_max_len_px:
|
||||
global_max_len_px = length
|
||||
best_x_rel = x
|
||||
best_y_start = y_s
|
||||
best_y_end = y_e
|
||||
|
||||
if global_max_len_px == 0:
|
||||
return None
|
||||
|
||||
thickness_mm = global_max_len_px * PIXEL_TO_MM
|
||||
real_x = roi_start + best_x_rel
|
||||
|
||||
print(f"📏 Measurement: {thickness_mm:.2f}mm ({global_max_len_px}px) at x={real_x}")
|
||||
|
||||
return {
|
||||
'thickness_px': int(global_max_len_px),
|
||||
'thickness_mm': float(round(thickness_mm, 2)),
|
||||
'x': int(real_x),
|
||||
'y_start': int(best_y_start),
|
||||
'y_end': int(best_y_end),
|
||||
'roi_start': int(roi_start),
|
||||
'roi_end': int(roi_end),
|
||||
'bbox': {'x': int(x_all), 'y': int(y_all), 'w': int(w_all), 'h': int(h_all)}
|
||||
}
|
||||
|
||||
def analyze_inflammation_severity(masks, image_size):
|
||||
if not masks:
|
||||
return None
|
||||
|
||||
width, height = image_size
|
||||
total_pixels = width * height
|
||||
|
||||
effusion_mask = masks.get('effusion', np.zeros((height, width), dtype=np.uint8))
|
||||
effusion_pixels = int(np.sum(effusion_mask))
|
||||
effusion_ratio = (effusion_pixels / total_pixels) * 100
|
||||
|
||||
effusion_thickness = 0
|
||||
if effusion_pixels > 0:
|
||||
rows_with_effusion = np.any(effusion_mask > 0, axis=1)
|
||||
if np.any(rows_with_effusion):
|
||||
effusion_thickness = int(np.sum(rows_with_effusion))
|
||||
|
||||
synovium_mask = masks.get('synovium', np.zeros((height, width), dtype=np.uint8))
|
||||
synovium_pixels = int(np.sum(synovium_mask))
|
||||
synovium_ratio = (synovium_pixels / total_pixels) * 100
|
||||
|
||||
effusion_score = min(effusion_thickness / height * 100, 100)
|
||||
synovium_score = synovium_ratio
|
||||
combined_score = (effusion_score * 0.6 + synovium_score * 0.4)
|
||||
|
||||
if combined_score > 15:
|
||||
level, severity, color = 3, "Nặng", "#dc3545"
|
||||
description = f"Dịch khớp dày ({effusion_thickness}px), màng hoạt dịch tăng sinh rõ"
|
||||
elif combined_score >= 8:
|
||||
level, severity, color = 2, "Trung bình", "#fd7e14"
|
||||
description = f"Dịch khớp trung bình ({effusion_thickness}px), màng hoạt dịch tăng sinh vừa"
|
||||
elif combined_score >= 3:
|
||||
level, severity, color = 1, "Nhẹ", "#ffc107"
|
||||
description = f"Dịch khớp mỏng ({effusion_thickness}px), màng hoạt dịch tăng sinh nhẹ"
|
||||
else:
|
||||
level, severity, color = 0, "Rất nhẹ", "#28a745"
|
||||
description = "Lượng dịch và màng hoạt dịch trong giới hạn bình thường"
|
||||
|
||||
return {
|
||||
'level': int(level),
|
||||
'severity': severity,
|
||||
'color': color,
|
||||
'description': description,
|
||||
'effusion': {'pixels': effusion_pixels, 'ratio': float(round(effusion_ratio, 2)), 'thickness': effusion_thickness},
|
||||
'synovium': {'pixels': synovium_pixels, 'ratio': float(round(synovium_ratio, 2))},
|
||||
'combined_score': float(round(combined_score, 2))
|
||||
}
|
||||
|
||||
def create_segmentation_overlay(image_pil, masks, measurement=None, angle_type='sup'):
|
||||
if masks is None:
|
||||
return image_pil
|
||||
|
||||
color_map = COLOR_MAP_SUP if angle_type == 'sup' else COLOR_MAP_POST
|
||||
|
||||
img_array = np.array(image_pil)
|
||||
overlay = img_array.copy()
|
||||
|
||||
for class_name, mask in masks.items():
|
||||
if class_name in color_map and np.sum(mask) > 0:
|
||||
color = color_map[class_name]['color']
|
||||
for i in range(3):
|
||||
overlay[:, :, i] = np.where(mask > 0,
|
||||
(overlay[:, :, i] * 0.6 + color[i] * 0.4).astype(np.uint8),
|
||||
overlay[:, :, i])
|
||||
|
||||
overlay_pil = Image.fromarray(overlay)
|
||||
draw = ImageDraw.Draw(overlay_pil)
|
||||
|
||||
for class_name in ['effusion', 'synovium']:
|
||||
mask = masks.get(class_name)
|
||||
if mask is not None and np.sum(mask) > 0:
|
||||
mask_uint8 = (mask * 255).astype(np.uint8)
|
||||
contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
|
||||
for contour in contours:
|
||||
points = contour.reshape(-1, 2).tolist()
|
||||
if len(points) > 2:
|
||||
points = [(int(p[0]), int(p[1])) for p in points]
|
||||
draw.line(points + [points[0]], fill=(255, 255, 255), width=3)
|
||||
|
||||
if measurement and angle_type == 'sup':
|
||||
x = measurement['x']
|
||||
y_start = measurement['y_start']
|
||||
y_end = measurement['y_end']
|
||||
thickness_mm = measurement['thickness_mm']
|
||||
roi_start = measurement['roi_start']
|
||||
roi_end = measurement['roi_end']
|
||||
bbox = measurement['bbox']
|
||||
|
||||
draw.rectangle(
|
||||
[bbox['x'], bbox['y'], bbox['x'] + bbox['w'], bbox['y'] + bbox['h']],
|
||||
outline=(0, 255, 0), # Chuyển sang xanh lá cho dễ nhìn
|
||||
width=3 # Tăng độ dày khung
|
||||
)
|
||||
|
||||
h = image_pil.size[1]
|
||||
draw.line([(roi_start, 0), (roi_start, h)], fill=(0, 255, 255), width=2)
|
||||
draw.line([(roi_end, 0), (roi_end, h)], fill=(0, 255, 255), width=2)
|
||||
|
||||
draw.line([(x, y_start), (x, y_end)], fill=(255, 0, 0), width=4)
|
||||
|
||||
radius = 4
|
||||
draw.ellipse([x-radius, y_start-radius, x+radius, y_start+radius],
|
||||
fill=(0, 255, 0), outline=(255, 255, 255), width=2)
|
||||
draw.ellipse([x-radius, y_end-radius, x+radius, y_end+radius],
|
||||
fill=(0, 255, 0), outline=(255, 255, 255), width=2)
|
||||
|
||||
text = f"{thickness_mm:.2f} mm"
|
||||
|
||||
try:
|
||||
from PIL import ImageFont
|
||||
font = ImageFont.load_default()
|
||||
bbox_text = draw.textbbox((0, 0), text, font=font)
|
||||
text_w = bbox_text[2] - bbox_text[0]
|
||||
text_h = bbox_text[3] - bbox_text[1]
|
||||
except:
|
||||
text_w, text_h = 100, 20
|
||||
|
||||
text_x = x + 8
|
||||
text_y = y_start - text_h - 8
|
||||
|
||||
draw.rectangle(
|
||||
[text_x - 2, text_y - 2, text_x + text_w + 2, text_y + text_h + 2],
|
||||
fill=(0, 0, 0)
|
||||
)
|
||||
draw.text((text_x, text_y), text, fill=(255, 0, 0))
|
||||
|
||||
return overlay_pil
|
||||
|
||||
def apply_clahe(image_pil):
|
||||
"""Áp dụng thuật toán CLAHE để tăng độ tương phản. Phục vụ cả hiển thị và làm đầu vào AI."""
|
||||
# Chuyển từ PIL sang OpenCV (numpy array)
|
||||
img_array = np.array(image_pil)
|
||||
# Chuyển sang thang độ xám (Gray) để xử lý CLAHE
|
||||
gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
|
||||
|
||||
# Tạo đối tượng CLAHE
|
||||
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
|
||||
enhanced_gray = clahe.apply(gray)
|
||||
|
||||
# Chuyển ngược lại sang RGB (3 kênh) để tương thích với các models
|
||||
enhanced_rgb = cv2.cvtColor(enhanced_gray, cv2.COLOR_GRAY2RGB)
|
||||
return Image.fromarray(enhanced_rgb)
|
||||
|
||||
# Mount thư mục static (CSS, JS)
|
||||
app.mount("/css", StaticFiles(directory="templates/css"), name="css")
|
||||
app.mount("/js", StaticFiles(directory="templates/js"), name="js")
|
||||
|
||||
@app.get("/")
|
||||
async def read_index():
|
||||
html_file = Path(TEMPLATES_FOLDER) / "index.html"
|
||||
if html_file.exists():
|
||||
return FileResponse(html_file)
|
||||
return JSONResponse({"error": "Template not found"})
|
||||
|
||||
@app.post("/api/analyze")
|
||||
async def analyze_image(
|
||||
image: UploadFile = File(...),
|
||||
angle_model: str = Query("convnext"),
|
||||
inflammation_model: str = Query("efficientnet_b0"),
|
||||
segment_model_sup: str = Query("deeplabv3"),
|
||||
segment_model_post: str = Query("deeplabv3_resnet101")
|
||||
):
|
||||
try:
|
||||
print(f"\n{'='*60}")
|
||||
print(f"📊 NEW REQUEST")
|
||||
print(f"Models: angle={angle_model}, inflam={inflammation_model}")
|
||||
print(f" seg_sup={segment_model_sup}, seg_post={segment_model_post}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
contents = await image.read()
|
||||
image_pil = Image.open(io.BytesIO(contents)).convert('RGB')
|
||||
|
||||
# Tạo ảnh tăng cường độ tương phản (Enhanced) cho mục đích hiển thị
|
||||
enhanced_pil = apply_clahe(image_pil)
|
||||
buffered_en = io.BytesIO()
|
||||
enhanced_pil.save(buffered_en, format="PNG")
|
||||
enhanced_str = base64.b64encode(buffered_en.getvalue()).decode()
|
||||
|
||||
result = {
|
||||
'success': True,
|
||||
'filename': image.filename,
|
||||
'images': {
|
||||
'enhanced': f"data:image/png;base64,{enhanced_str}"
|
||||
},
|
||||
'models_used': {
|
||||
'angle': angle_model,
|
||||
'inflammation': inflammation_model,
|
||||
'segmentation_sup': segment_model_sup,
|
||||
'segmentation_post': segment_model_post
|
||||
}
|
||||
}
|
||||
|
||||
angle_clf = load_angle_model(angle_model)
|
||||
angle, angle_conf = predict_angle(angle_clf, image_pil)
|
||||
result['angle'] = {'class': angle, 'confidence': angle_conf}
|
||||
print(f"✅ Angle: {angle} ({angle_conf}%)")
|
||||
|
||||
if 'post-trans' in angle.lower():
|
||||
print(f"🔍 POST-TRANS pipeline")
|
||||
|
||||
inflam_model = load_inflammation_model()
|
||||
has_inflammation, inflam_conf = predict_inflammation(inflam_model, image_pil)
|
||||
result['inflammation'] = {'detected': has_inflammation, 'confidence': inflam_conf}
|
||||
print(f"✅ Inflammation: {has_inflammation} ({inflam_conf}%)")
|
||||
|
||||
if has_inflammation:
|
||||
seg_model = load_segmentation_model_post(segment_model_post)
|
||||
preds, masks = segment_image(seg_model, image_pil, segment_model_post, 'post')
|
||||
|
||||
if masks:
|
||||
segmented_img = create_segmentation_overlay(image_pil, masks, None, 'post')
|
||||
# Convert to base64
|
||||
buffered = io.BytesIO()
|
||||
segmented_img.save(buffered, format="PNG")
|
||||
img_str = base64.b64encode(buffered.getvalue()).decode()
|
||||
|
||||
result['images']['segmented'] = f"data:image/png;base64,{img_str}"
|
||||
|
||||
detected_classes = [k for k, v in masks.items() if np.sum(v) > 0]
|
||||
color_legend = []
|
||||
for class_name in detected_classes:
|
||||
if class_name in COLOR_MAP_POST:
|
||||
color_legend.append({
|
||||
'name': COLOR_MAP_POST[class_name]['name'],
|
||||
'color': COLOR_MAP_POST[class_name]['color'],
|
||||
'key': class_name
|
||||
})
|
||||
|
||||
result['segmentation'] = {
|
||||
'performed': True,
|
||||
'classes_detected': detected_classes,
|
||||
'color_legend': color_legend,
|
||||
'angle_type': 'post'
|
||||
}
|
||||
|
||||
print(f"✅ Segmentation POST completed")
|
||||
else:
|
||||
print(f"ℹ️ No inflammation detected - skipping segmentation POST")
|
||||
|
||||
elif 'sup-up-long' in angle.lower():
|
||||
print(f"🔍 SUPRAPAT pipeline")
|
||||
|
||||
inflam_model = load_inflammation_model()
|
||||
has_inflammation, inflam_conf = predict_inflammation(inflam_model, image_pil)
|
||||
result['inflammation'] = {'detected': has_inflammation, 'confidence': inflam_conf}
|
||||
print(f"✅ Inflammation: {has_inflammation} ({inflam_conf}%)")
|
||||
|
||||
if has_inflammation:
|
||||
seg_model = load_segmentation_model_sup(segment_model_sup)
|
||||
preds, masks = segment_image(seg_model, image_pil, segment_model_sup, 'sup')
|
||||
|
||||
if masks:
|
||||
measurement = measure_thickness_new(masks, image_pil.size)
|
||||
|
||||
segmented_img = create_segmentation_overlay(image_pil, masks, measurement, 'sup')
|
||||
# Convert to base64
|
||||
buffered = io.BytesIO()
|
||||
segmented_img.save(buffered, format="PNG")
|
||||
img_str = base64.b64encode(buffered.getvalue()).decode()
|
||||
|
||||
result['images']['segmented'] = f"data:image/png;base64,{img_str}"
|
||||
|
||||
if measurement:
|
||||
result['measurement'] = {
|
||||
'thickness_mm': measurement['thickness_mm'],
|
||||
'thickness_px': measurement['thickness_px'],
|
||||
'location_x': measurement['x'],
|
||||
'y_start': measurement['y_start'],
|
||||
'y_end': measurement['y_end']
|
||||
}
|
||||
print(f"✅ Measurement: {measurement['thickness_mm']:.2f}mm at x={measurement['x']}")
|
||||
|
||||
detected_classes = [k for k, v in masks.items() if np.sum(v) > 0]
|
||||
color_legend = []
|
||||
for class_name in detected_classes:
|
||||
if class_name in COLOR_MAP_SUP:
|
||||
color_legend.append({
|
||||
'name': COLOR_MAP_SUP[class_name]['name'],
|
||||
'color': COLOR_MAP_SUP[class_name]['color'],
|
||||
'key': class_name
|
||||
})
|
||||
|
||||
result['segmentation'] = {
|
||||
'performed': True,
|
||||
'classes_detected': detected_classes,
|
||||
'color_legend': color_legend,
|
||||
'angle_type': 'sup'
|
||||
}
|
||||
|
||||
severity = analyze_inflammation_severity(masks, image_pil.size)
|
||||
if severity:
|
||||
result['severity'] = severity
|
||||
print(f"✅ Severity: {severity['severity']}")
|
||||
|
||||
print(f"✅ Segmentation SUP completed")
|
||||
else:
|
||||
print(f"ℹ️ No inflammation detected - skipping segmentation SUP")
|
||||
|
||||
else:
|
||||
print(f"ℹ️ Other angle - only angle classification")
|
||||
|
||||
print(f"{'='*60}\n")
|
||||
return JSONResponse(result)
|
||||
|
||||
except Exception as e:
|
||||
import traceback
|
||||
print(f"❌ Error: {e}")
|
||||
print(traceback.format_exc())
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get("/api/health")
|
||||
async def health_check():
|
||||
return JSONResponse({'status': 'healthy'})
|
||||
|
||||
def sanitize_name(name):
|
||||
# Loại bỏ ký tự không hợp lệ cho folder trên Windows
|
||||
if not name: return "unknown"
|
||||
# Thay thế ký tự lạ bằng dấu gạch dưới
|
||||
clean = re.sub(r'[\\/*?:"<>|]', "_", name)
|
||||
# Loại bỏ dấu cách thừa
|
||||
clean = clean.strip().replace(" ", "_")
|
||||
return clean
|
||||
|
||||
class SaveDataRequest(BaseModel):
|
||||
patient_info: dict
|
||||
analysis_result: dict
|
||||
images: dict
|
||||
|
||||
@app.post("/api/save")
|
||||
async def save_patient_data(data: SaveDataRequest):
|
||||
try:
|
||||
p = data.patient_info
|
||||
res = data.analysis_result
|
||||
imgs = data.images
|
||||
|
||||
# 1. Tạo thư mục theo mã bệnh nhân (nhóm chính)
|
||||
patient_id = sanitize_name(p.get('id', 'unknown'))
|
||||
patient_name = sanitize_name(p.get('name', 'no_name'))
|
||||
|
||||
# Thư mục chính của bệnh nhân
|
||||
patient_folder = f"{patient_id}_{patient_name}"
|
||||
|
||||
# Thư mục con theo thời gian hiện tại
|
||||
timestamp_folder = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
|
||||
# Tổng hợp đường dẫn: patients/ID_Name/TIMESTAMP/
|
||||
target_dir = os.path.join("patients", patient_folder, timestamp_folder)
|
||||
os.makedirs(target_dir, exist_ok=True)
|
||||
|
||||
# 2. Lưu info.txt
|
||||
info_path = os.path.join(target_dir, "info.txt")
|
||||
with open(info_path, "w", encoding="utf-8") as f:
|
||||
f.write(f"--- THÔNG TIN BỆNH NHÂN ---\n")
|
||||
f.write(f"Mã BN: {patient_id}\n")
|
||||
f.write(f"Họ tên: {patient_name}\n")
|
||||
f.write(f"Giới tính: {p.get('gender')}\n")
|
||||
f.write(f"Tuổi: {p.get('age')}\n")
|
||||
f.write(f"Chẩn đoán BS: {p.get('diagnosis')}\n\n")
|
||||
|
||||
f.write(f"--- KẾT QUẢ PHÂN TÍCH AI ---\n")
|
||||
f.write(f"Góc chụp: {res.get('angle', {}).get('class')} ({res.get('angle', {}).get('confidence')}%)\n")
|
||||
|
||||
if 'inflammation' in res:
|
||||
infl = res['inflammation']
|
||||
f.write(f"Viêm nhiễm: {'Có' if infl['detected'] else 'Không'} ({infl['confidence']}%)\n")
|
||||
|
||||
if 'measurement' in res:
|
||||
m = res['measurement']
|
||||
f.write(f"Độ dày màng: {m['thickness_mm']} mm ({m['thickness_px']} px)\n")
|
||||
f.write(f"Vị trí x: {m['location_x']}\n")
|
||||
|
||||
if 'severity' in res:
|
||||
s = res['severity']
|
||||
f.write(f"Mức độ: {s['severity']}\n")
|
||||
f.write(f"Mô tả: {s['description']}\n")
|
||||
|
||||
# 3. Lưu ảnh
|
||||
def save_base64_img(b64_str, filename):
|
||||
if not b64_str: return
|
||||
# Remove header if present
|
||||
if "," in b64_str:
|
||||
b64_str = b64_str.split(",")[1]
|
||||
|
||||
img_data = base64.b64decode(b64_str)
|
||||
with open(os.path.join(target_dir, filename), "wb") as f:
|
||||
f.write(img_data)
|
||||
|
||||
save_base64_img(imgs.get('original'), "original.png")
|
||||
save_base64_img(imgs.get('segmented'), "segmented.png")
|
||||
|
||||
# 4. Tự động lưu PDF báo cáo
|
||||
try:
|
||||
pdf_bytes = generate_medical_report(p, res, imgs)
|
||||
pdf_path = os.path.join(target_dir, "report.pdf")
|
||||
with open(pdf_path, "wb") as f:
|
||||
f.write(bytes(pdf_bytes))
|
||||
print(f"📄 Report PDF saved to: {pdf_path}")
|
||||
except Exception as pdf_err:
|
||||
print(f"⚠️ Warning: Could not auto-save PDF: {pdf_err}")
|
||||
|
||||
print(f"✅ Data saved for patient: {patient_id}")
|
||||
return {"success": True, "folder": target_dir}
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Save Error: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post("/api/export-pdf")
|
||||
async def export_patient_pdf(data: SaveDataRequest):
|
||||
try:
|
||||
pdf_bytes = generate_medical_report(
|
||||
data.patient_info,
|
||||
data.analysis_result,
|
||||
data.images
|
||||
)
|
||||
|
||||
filename = f"Phieu_Kham_{sanitize_name(data.patient_info.get('id', 'unknown'))}.pdf"
|
||||
|
||||
return Response(
|
||||
content=bytes(pdf_bytes),
|
||||
media_type="application/pdf",
|
||||
headers={"Content-Disposition": f"attachment; filename={filename}"}
|
||||
)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
print(f"❌ PDF Export Error: {e}")
|
||||
print(traceback.format_exc())
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
if __name__ == '__main__':
|
||||
print("Medical Image Analysis Server")
|
||||
print(f"URL: http://127.0.0.1:8000")
|
||||
print(f"Device: {device}")
|
||||
uvicorn.run(app, host="127.0.0.1", port=8000)
|
||||
Reference in New Issue
Block a user