George-waheed/fracture
0
1# ================================
2# Imports
3# ================================
4import os
5import torch
6import torch.nn as nn
7import torchvision.models as models
8
9from torchvision.models.detection import fasterrcnn_resnet50_fpn
10from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
11from torchvision import transforms
12
13import cv2
14
15import matplotlib
16matplotlib.use("Agg")
17
18import matplotlib.pyplot as plt
19import matplotlib.patches as patches
20
21from PIL import Image
22import numpy as np
23
24# ================================
25# Device
26# ================================
27DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
28
29# ================================
30# Config
31# ================================
32CLF_MODEL_PATH = "best_model.pth"
33DET_MODEL_PATH = "final_fracture_model_v2.pth"
34
35DET_CONF_THRESHOLD = 0.3
36
37# ================================
38# Treatment Recommendations
39# ================================
40# ─────────────────────────────────────────────────────────────
41# 🔥 TREATMENT — 3 حالات × 3 مراحل + Red Flags
42# ─────────────────────────────────────────────────────────────
43TREATMENT = {
44
45 "fractured": {
46 "title": "🦴 FRACTURE DETECTED",
47 "urgency": "🔴 URGENT — Seek medical attention within 24 hours",
48
49 "phase_1_emergency": {
50 "label": "🚨 Phase 1 — Immediate First Aid",
51 "steps": [
52 "🛑 Do NOT move or bend the injured limb",
53 "🩹 If bleeding present: apply gentle pressure with clean cloth",
54 "🧊 Apply ice pack (wrapped in cloth) — 20 min ON / 20 min OFF",
55 "🚑 Call emergency services or go to ER immediately",
56 "📵 Do NOT attempt to realign the bone yourself",
57 ]
58 },
59
60 "phase_2_hospital": {
61 "label": "🏥 Phase 2 — Hospital Care",
62 "steps": [
63 "🔬 X-ray / CT scan for confirmation and fracture type",
64 "💉 Pain management (IV or oral as per doctor)",
65 "🦺 Splint, cast, or surgical fixation based on severity",
66 "🩺 Orthopedic surgeon consultation if displaced fracture",
67 ]
68 },
69
70 "phase_3_recovery": {
71 "label": "🗓️ Phase 3 — Recovery",
72 "steps": [
73 "📅 Follow-up appointment in 1–2 weeks",
74 "🏋️ Physical therapy after immobilization period",
75 "🥛 Calcium (1000–1200 mg/day) + Vitamin D supplements",
76 "🚫 No weight-bearing on injured limb until doctor approves",
77 "😴 Elevate the limb to reduce swelling during rest",
78 ]
79 },
80
81 "red_flags": {
82 "label": "⚠️ Red Flags — Go to ER Immediately if:",
83 "steps": [
84 "🫀 Numbness, tingling, or coldness below injury → nerve/vessel damage",
85 "🩸 Bone visible through skin → open fracture, high infection risk",
86 "😵 Dizziness, pale skin, rapid heartbeat → shock symptoms",
87 "🌡️ Fever after fracture → possible bone infection (osteomyelitis)",
88 "💜 Severe swelling / visible deformity → possible displaced fracture",
89 ]
90 },
91 },
92
93 "not fractured": {
94 "title": "✅ NO FRACTURE DETECTED",
95 "urgency": "🟢 No emergency — Monitor symptoms",
96
97 "phase_1_emergency": {
98 "label": "🧊 Immediate Care",
99 "steps": [
100 "🧊 Apply cold compress for 15–20 min to reduce swelling",
101 "🛋️ Rest the injured area — avoid strenuous activity",
102 "🩹 Compress with elastic bandage if swelling present",
103 "⬆️ Elevate the limb above heart level when resting",
104 ]
105 },
106
107 "phase_2_hospital": {
108 "label": "💊 Short-Term Management",
109 "steps": [
110 "💊 Anti-inflammatory medication (ibuprofen/naproxen) if needed",
111 "🚶 Gradually return to normal activities over 2–3 days",
112 "🧘 Light stretching after pain subsides",
113 ]
114 },
115
116 "phase_3_recovery": {
117 "label": "🔁 When to Re-Evaluate",
118 "steps": [
119 "📅 Re-check if pain persists more than 3–5 days",
120 "🏥 Visit a doctor if swelling worsens or new symptoms appear",
121 "🧪 Consider repeat imaging if pain is severe despite no fracture",
122 ]
123 },
124
125 "red_flags": {
126 "label": "⚠️ Red Flags — Visit a Doctor if:",
127 "steps": [
128 "📈 Pain is getting worse instead of better after 48 hours",
129 "🦵 Unable to bear weight at all → possible stress fracture",
130 "🌡️ Fever or redness around injury → possible infection",
131 "😶🌫️ Persistent numbness → possible nerve involvement",
132 ]
133 },
134 },
135
136 "uncertain": {
137 "title": "⚠️ POSSIBLE FRACTURE — UNCERTAIN",
138 "urgency": "🟠 CAUTION — Medical review recommended within 24–48 hours",
139
140 "phase_1_emergency": {
141 "label": "🛡️ Immediate Precaution",
142 "steps": [
143 "📌 AI models gave conflicting results — fracture cannot be ruled out",
144 "🛑 Immobilize the affected area as a precaution",
145 "🧊 Apply cold compress (wrapped ice) — 20 min ON / 20 min OFF",
146 "🚫 Avoid putting any weight or pressure on the area",
147 ]
148 },
149
150 "phase_2_hospital": {
151 "label": "🏥 Medical Evaluation",
152 "steps": [
153 "🏥 Visit an orthopedic doctor or urgent care within 24–48 hours",
154 "🧪 Additional imaging required: X-ray or CT scan",
155 "💊 Pain relief medication as prescribed by physician",
156 "🩺 Physical examination to rule out occult fracture",
157 ]
158 },
159
160 "phase_3_recovery": {
161 "label": "🗓️ Follow-Up",
162 "steps": [
163 "📅 Do not ignore persistent pain — follow up if symptoms continue",
164 "🥛 Maintain calcium & vitamin D intake for bone health",
165 "📋 Keep record of symptoms to share with your doctor",
166 ]
167 },
168
169 "red_flags": {
170 "label": "⚠️ Red Flags — Go to ER Immediately if:",
171 "steps": [
172 "🫀 Numbness or coldness below injury site",
173 "🩸 Open wound or bone visible through skin",
174 "😵 Signs of shock: dizziness, pale skin, rapid heartbeat",
175 "💜 Rapidly increasing swelling or visible deformity",
176 ]
177 },
178 }
179}
180
181# ================================
182# Load Classification Model
183# ================================
184def load_classifier(path):
185
186 model = models.efficientnet_b0(weights=None)
187
188 in_features = model.classifier[1].in_features
189
190 model.classifier = nn.Sequential(
191 nn.Dropout(0.3),
192 nn.Linear(in_features, 2),
193 )
194
195 model.load_state_dict(
196 torch.load(path, map_location=DEVICE)
197 )
198
199 model.to(DEVICE)
200 model.eval()
201
202 return model
203
204
205clf_model = load_classifier(CLF_MODEL_PATH)
206
207# ================================
208# Load Detection Model
209# ================================
210def load_detector(path):
211
212 model = fasterrcnn_resnet50_fpn(weights=None)
213
214 in_features = model.roi_heads.box_predictor.cls_score.in_features
215
216 model.roi_heads.box_predictor = FastRCNNPredictor(
217 in_features,
218 2
219 )
220
221 model.load_state_dict(
222 torch.load(path, map_location=DEVICE)
223 )
224
225 model.to(DEVICE)
226 model.eval()
227
228 return model
229
230
231det_model = load_detector(DET_MODEL_PATH)
232
233# ================================
234# Classification Transform
235# ================================
236clf_transform = transforms.Compose([
237 transforms.Resize((224, 224)),
238 transforms.ToTensor(),
239 transforms.Normalize(
240 [0.485, 0.456, 0.406],
241 [0.229, 0.224, 0.225]
242 )
243])
244
245CLF_CLASSES = ['fractured', 'not fractured']
246
247# ================================
248# Main Pipeline
249# ================================
250def run_pipeline(img_path):
251
252 # ============================
253 # Read Image
254 # ============================
255 pil_img = Image.open(img_path).convert("RGB")
256
257 cv_img = cv2.cvtColor(
258 np.array(pil_img),
259 cv2.COLOR_RGB2BGR
260 )
261
262 h, w = cv_img.shape[:2]
263
264 # ============================
265 # Classification
266 # ============================
267 tensor = clf_transform(pil_img) \
268 .unsqueeze(0) \
269 .to(DEVICE)
270
271 with torch.no_grad():
272
273 logits = clf_model(tensor)
274
275 probs = torch.softmax(logits, dim=1)
276
277 conf, idx = torch.max(probs, 1)
278
279 clf_class = CLF_CLASSES[idx.item()]
280 clf_conf = conf.item()
281
282 # ============================
283 # Detection
284 # ============================
285 detected_boxes = []
286
287 img_resized = cv2.resize(cv_img, (800, 800))
288
289 img_rgb = cv2.cvtColor(
290 img_resized,
291 cv2.COLOR_BGR2RGB
292 )
293
294 det_tensor = torch.tensor(
295 img_rgb / 255.,
296 dtype=torch.float32
297 ).permute(2,0,1).unsqueeze(0).to(DEVICE)
298
299 with torch.no_grad():
300
301 output = det_model(det_tensor)[0]
302
303 for box, score, label in zip(
304 output['boxes'],
305 output['scores'],
306 output['labels']
307 ):
308
309 if score.item() >= DET_CONF_THRESHOLD and label.item() == 1:
310
311 x1, y1, x2, y2 = box.cpu().numpy()
312
313 x1 = int(x1 * w / 800)
314 y1 = int(y1 * h / 800)
315 x2 = int(x2 * w / 800)
316 y2 = int(y2 * h / 800)
317
318 detected_boxes.append(
319 (x1, y1, x2, y2, score.item())
320 )
321
322 # ============================
323 # Final Decision
324 # ============================
325 if clf_class == "fractured" and len(detected_boxes) > 0:
326
327 final_decision = "fractured"
328
329 elif clf_class == "not fractured" and len(detected_boxes) == 0:
330
331 final_decision = "not fractured"
332
333 else:
334
335 final_decision = "uncertain"
336
337 # ============================
338 # Visualization
339 # ============================
340 os.makedirs("outputs", exist_ok=True)
341
342 fig, ax = plt.subplots(1, 1, figsize=(5,5))
343
344 ax.imshow(
345 cv2.cvtColor(
346 cv_img,
347 cv2.COLOR_BGR2RGB
348 )
349 )
350
351 for (x1, y1, x2, y2, sc) in detected_boxes:
352
353 rect = patches.Rectangle(
354 (x1, y1),
355 x2 - x1,
356 y2 - y1,
357 linewidth=2,
358 edgecolor='red',
359 facecolor='none'
360 )
361
362 ax.add_patch(rect)
363
364 ax.text(
365 x1,
366 y1 - 5,
367 f"{sc:.2f}",
368 color='red'
369 )
370
371 ax.set_title(final_decision.upper())
372
373 ax.axis('off')
374
375 # ============================
376 # Save Output Image
377 # ============================
378 output_path = f"outputs/{final_decision}_result.jpg"
379
380 plt.savefig(
381 output_path,
382 bbox_inches='tight'
383 )
384
385 plt.close()
386
387 # ============================
388 # API Response
389 # ============================
390 return {
391
392 "decision": final_decision,
393
394 "confidence": round(clf_conf, 3),
395
396 "detections": len(detected_boxes),
397
398 "recommendation": TREATMENT[final_decision],
399
400 "image_url": f"/outputs/{final_decision}_result.jpg"
401 }