File size: 3,847 Bytes
a326b94 3bb1400 88fb5fa 0000f4a 88fb5fa 005d8cf 88fb5fa 3bb1400 88fb5fa 0000f4a 88fb5fa 0000f4a 88fb5fa 3bb1400 88fb5fa 3bb1400 912e146 3bb1400 6ea5ee2 3bb1400 912e146 88fb5fa 6ea5ee2 3bb1400 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
import streamlit as st
from transformers import pipeline
from PIL import Image, ImageDraw
import torch
st.set_page_config(page_title="Multi-Model Fracture Detection", layout="wide")
@st.cache_resource
def load_models():
models = {
"D3STRON (Object Detection)": pipeline("object-detection", model="D3STRON/bone-fracture-detr"),
"Heem2 (Classification)": pipeline("image-classification", model="Heem2/bone-fracture-detection-using-xray"),
"Akhileshav8 (Classification)": pipeline("image-classification", model="akhileshav8/image_classification_for_fracture"),
"Nandodeomkar (Classification)": pipeline("image-classification", model="nandodeomkar/autotrain-fracture-detection-using-google-vit-base-patch-16-54382127388"),
"Anirban22 (Object Detection)": pipeline("object-detection", model="anirban22/detr-resnet-50-med_fracture")
}
return models
def draw_boxes(image, predictions):
draw = ImageDraw.Draw(image)
for pred in predictions:
box = pred['box']
label = f"{pred['label']} ({pred['score']:.2%})"
draw.rectangle(
[(box['xmin'], box['ymin']), (box['xmax'], box['ymax'])],
outline="red",
width=3
)
text_bbox = draw.textbbox((box['xmin'], box['ymin']), label)
draw.rectangle(text_bbox, fill="red")
draw.text((box['xmin'], box['ymin']), label, fill="white")
return image
def process_classification(model, image, conf_threshold):
predictions = model(image)
results = []
for pred in predictions:
if pred['score'] >= conf_threshold:
results.append(f"{pred['label']}: {pred['score']:.2%}")
return results
def process_detection(model, image, conf_threshold):
predictions = model(image)
return [pred for pred in predictions if pred['score'] >= conf_threshold]
def main():
st.title("🦴 Multi-Model Fracture Detection")
models = load_models()
uploaded_file = st.file_uploader("Upload X-ray image", type=['png', 'jpg', 'jpeg'])
conf_threshold = st.slider(
"Confidence threshold",
min_value=0.0,
max_value=1.0,
value=0.3,
step=0.01
)
if uploaded_file:
image = Image.open(uploaded_file)
max_size = (400, 400)
image.thumbnail(max_size, Image.Resampling.LANCZOS)
st.image(image, caption="Original Image", width=400)
col1, col2 = st.columns(2)
with col1:
st.subheader("Classification Models")
for name, model in models.items():
if "Classification" in name:
st.write(f"**{name}**")
with st.spinner(f"Running {name}..."):
results = process_classification(model, image, conf_threshold)
for result in results:
st.write(f"• {result}")
with col2:
st.subheader("Object Detection Models")
for name, model in models.items():
if "Object Detection" in name:
st.write(f"**{name}**")
with st.spinner(f"Running {name}..."):
detections = process_detection(model, image, conf_threshold)
if detections:
result_image = image.copy()
result_image = draw_boxes(result_image, detections)
st.image(result_image, caption=f"Results from {name}")
for det in detections:
st.write(f"• {det['label']}: {det['score']:.2%}")
else:
st.write("No detections above threshold")
if __name__ == "__main__":
main() |