Poc1-Proof of Concept verison 1 #3
@@ -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)
|
||||||
|
|||||||
Reference in New Issue
Block a user