This commit is contained in:
DatTT127
2026-06-24 21:47:15 +07:00
parent 8de0c5e844
commit fed5f277f4

View File

@@ -84,101 +84,101 @@ app.add_middleware(
# ============ MODEL LOADING ============ # ============ MODEL LOADING ============
def load_angle_model(model_name): # def load_angle_model(model_name):
print(f"📄 Loading angle model: {model_name}") # print(f"📄 Loading angle model: {model_name}")
if model_name == "convnext": # if model_name == "convnext":
model = models.convnext_tiny(weights=None) # model = models.convnext_tiny(weights=None)
model.classifier[2] = nn.Linear(model.classifier[2].in_features, 4) # 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 = 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()} # checkpoint = {k.replace('model.', ''): v for k, v in checkpoint.items()}
model.load_state_dict(checkpoint) # model.load_state_dict(checkpoint)
elif model_name == "densenet": # elif model_name == "densenet":
model = models.densenet121(weights=None) # model = models.densenet121(weights=None)
model.classifier = nn.Linear(model.classifier.in_features, 4) # model.classifier = nn.Linear(model.classifier.in_features, 4)
checkpoint = torch.load(f"models/best_densenet.pth", map_location=device, weights_only=False) # 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()} # checkpoint = {k.replace('model.', ''): v for k, v in checkpoint.items()}
model.load_state_dict(checkpoint) # model.load_state_dict(checkpoint)
elif model_name == "resnet50": # elif model_name == "resnet50":
model = models.resnet50(weights=None) # model = models.resnet50(weights=None)
model.fc = nn.Linear(model.fc.in_features, 4) # model.fc = nn.Linear(model.fc.in_features, 4)
checkpoint = torch.load(f"models/best_resnet50.pth", map_location=device, weights_only=False) # 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()} # checkpoint = {k.replace('model.', ''): v for k, v in checkpoint.items()}
model.load_state_dict(checkpoint) # model.load_state_dict(checkpoint)
elif model_name == "efficientnet_b2": # elif model_name == "efficientnet_b2":
model = models.efficientnet_b2(weights=None) # model = models.efficientnet_b2(weights=None)
model.classifier[1] = nn.Linear(model.classifier[1].in_features, 4) # 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 = 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()} # checkpoint = {k.replace('model.', ''): v for k, v in checkpoint.items()}
model.load_state_dict(checkpoint) # model.load_state_dict(checkpoint)
elif model_name == "swin": # elif model_name == "swin":
model = models.swin_v2_s(weights=None) # model = models.swin_v2_s(weights=None)
model.head = nn.Linear(model.head.in_features, 4) # 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 = 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()} # checkpoint = {k.replace('model.', ''): v for k, v in checkpoint.items()}
model.load_state_dict(checkpoint) # model.load_state_dict(checkpoint)
else: # else:
raise ValueError(f"Unknown angle model: {model_name}") # raise ValueError(f"Unknown angle model: {model_name}")
print(f"✅ Loaded: {model_name}") # print(f"✅ Loaded: {model_name}")
return model.to(device).eval() # return model.to(device).eval()
def load_inflammation_model(): # def load_inflammation_model():
print("📄 Loading inflammation model") # print("📄 Loading inflammation model")
model = models.efficientnet_b0(weights=None) # model = models.efficientnet_b0(weights=None)
model.classifier[1] = nn.Linear(model.classifier[1].in_features, 2) # 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)) # model.load_state_dict(torch.load("models/efficientnet_b0_ultrasound_2_class.pth", map_location=device, weights_only=False))
print("✅ Loaded inflammation model") # print("✅ Loaded inflammation model")
return model.to(device).eval() # return model.to(device).eval()
def load_segmentation_model_sup(model_name): # def load_segmentation_model_sup(model_name):
print(f"📄 Loading segmentation model SUP: {model_name}") # print(f"📄 Loading segmentation model SUP: {model_name}")
if model_name == "deeplabv3": # if model_name == "deeplabv3":
model = models.segmentation.deeplabv3_resnet50(weights=None) # model = models.segmentation.deeplabv3_resnet50(weights=None)
in_ch = model.classifier[-1].in_channels # in_ch = model.classifier[-1].in_channels
model.classifier = nn.Sequential( # model.classifier = nn.Sequential(
model.classifier[0], # model.classifier[0],
nn.Dropout(0.3), # nn.Dropout(0.3),
nn.Conv2d(in_ch, 7, kernel_size=1) # 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) # model.load_state_dict(torch.load("models/best_model_Deeplav3.pth", map_location=device, weights_only=False), strict=False)
elif model_name == "unet_resnet101": # elif model_name == "unet_resnet101":
try: # try:
from segmentation_models_pytorch import Unet # from segmentation_models_pytorch import Unet
model = Unet(encoder_name="resnet101", encoder_weights=None, classes=7) # 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)) # model.load_state_dict(torch.load("models/unet_resnet101.pth", map_location=device, weights_only=False))
except ImportError: # except ImportError:
raise ValueError("segmentation_models_pytorch not installed") # raise ValueError("segmentation_models_pytorch not installed")
elif model_name == "efficientfeedback": # elif model_name == "efficientfeedback":
model = EfficientFeedbackNetwork(in_channels=3, num_class=7) # model = EfficientFeedbackNetwork(in_channels=3, num_class=7)
model.load_state_dict(torch.load("models/efficientfeedback.pth", map_location=device, weights_only=False)) # model.load_state_dict(torch.load("models/efficientfeedback.pth", map_location=device, weights_only=False))
elif model_name == "unet3plus": # elif model_name == "unet3plus":
model = UNet3Plus_Attention(in_channels=3, num_classes=7) # 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)) # model.load_state_dict(torch.load("models/unet3plus_att.pth", map_location=device, weights_only=False))
else: # else:
raise ValueError(f"Unknown segmentation model: {model_name}") # raise ValueError(f"Unknown segmentation model: {model_name}")
print(f"✅ Loaded SUP: {model_name}") # print(f"✅ Loaded SUP: {model_name}")
return model.to(device).eval() # return model.to(device).eval()
def load_segmentation_model_post(model_name): # def load_segmentation_model_post(model_name):
print(f"📄 Loading segmentation model POST: {model_name}") # print(f"📄 Loading segmentation model POST: {model_name}")
if model_name == "deeplabv3_resnet101": # if model_name == "deeplabv3_resnet101":
model = models.segmentation.deeplabv3_resnet101(weights=None) # model = models.segmentation.deeplabv3_resnet101(weights=None)
in_ch = model.classifier[-1].in_channels # in_ch = model.classifier[-1].in_channels
model.classifier = nn.Sequential( # model.classifier = nn.Sequential(
model.classifier[0], # model.classifier[0],
nn.Dropout(0.3), # nn.Dropout(0.3),
nn.Conv2d(in_ch, 7, kernel_size=1) # 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) # model.load_state_dict(torch.load("models/best_model_deeplabv3_resnet101_seed_16.pth", map_location=device, weights_only=False), strict=False)
else: # else:
raise ValueError(f"Unknown post segmentation model: {model_name}") # raise ValueError(f"Unknown post segmentation model: {model_name}")
print(f"✅ Loaded POST: {model_name}") # print(f"✅ Loaded POST: {model_name}")
return model.to(device).eval() # return model.to(device).eval()
# Transforms # Transforms
angle_transform = transforms.Compose([ angle_transform = transforms.Compose([
@@ -212,7 +212,7 @@ def predict_angle(model, image_pil):
@torch.no_grad() @torch.no_grad()
def predict_inflammation(model, image_pil): def predict_inflammation(model, image_pil):
img_tensor = inflammation_transform(image_pil).unsqueeze(0).to(device) img_tensor = inflammation_transform(image_pil).unsqueeze(0).to(device)
output = model(img_tensor) output = model(img_tensor) # TRITON
probs = torch.softmax(output, dim=1) probs = torch.softmax(output, dim=1)
pred_class = torch.argmax(probs, dim=1).item() pred_class = torch.argmax(probs, dim=1).item()
confidence = probs[0][pred_class].item() confidence = probs[0][pred_class].item()
@@ -559,16 +559,17 @@ def apply_clahe(image_pil):
enhanced_rgb = cv2.cvtColor(enhanced_gray, cv2.COLOR_GRAY2RGB) enhanced_rgb = cv2.cvtColor(enhanced_gray, cv2.COLOR_GRAY2RGB)
return Image.fromarray(enhanced_rgb) return Image.fromarray(enhanced_rgb)
# Mount thư mục static (CSS, JS) # # Mount thư mục static (CSS, JS)
app.mount("/css", StaticFiles(directory="templates/css"), name="css") # app.mount("/css", StaticFiles(directory="templates/css"), name="css")
app.mount("/js", StaticFiles(directory="templates/js"), name="js") # app.mount("/js", StaticFiles(directory="templates/js"), name="js")
@app.get("/")
async def read_index(): # @app.get("/")
html_file = Path(TEMPLATES_FOLDER) / "index.html" # async def read_index():
if html_file.exists(): # html_file = Path(TEMPLATES_FOLDER) / "index.html"
return FileResponse(html_file) # if html_file.exists():
return JSONResponse({"error": "Template not found"}) # return FileResponse(html_file)
# return JSONResponse({"error": "Template not found"})
@app.post("/api/analyze") @app.post("/api/analyze")
async def analyze_image( async def analyze_image(
@@ -665,7 +666,7 @@ async def analyze_image(
if has_inflammation: if has_inflammation:
seg_model = load_segmentation_model_sup(segment_model_sup) seg_model = load_segmentation_model_sup(segment_model_sup)
preds, masks = segment_image(seg_model, image_pil, segment_model_sup, 'sup') _, masks = segment_image(seg_model, image_pil, segment_model_sup, 'sup')
if masks: if masks:
measurement = measure_thickness_new(masks, image_pil.size) measurement = measure_thickness_new(masks, image_pil.size)