Update app.py

app.py

@@ -3,20 +3,45 @@ from transformers import pipeline
 from PIL import Image, ImageDraw
 import torch
 
-st.set_page_config(page_title="Multi-Model Fracture Detection", layout="wide")
+st.set_page_config(
+    page_title="Fracture Detection",
+    layout="wide",
+    initial_sidebar_state="collapsed"
+)
+
+# Custom CSS
+st.markdown("""
+<style>
+    .main > div {
+        padding: 2rem;
+        background: #f8f9fa;
+        border-radius: 1rem;
+        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
+    }
+    .stButton button {
+        width: 100%;
+        border-radius: 0.5rem;
+    }
+    .uploadedFile {
+        border-radius: 0.5rem;
+    }
+    h1, h2, h3 {
+        color: #2c3e50;
+    }
+</style>
+""", unsafe_allow_html=True)
 
 @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 {
+        "D3STRON": pipeline("object-detection", model="D3STRON/bone-fracture-detr"),
+        "Heem2": pipeline("image-classification", model="Heem2/bone-fracture-detection-using-xray"),
+        "Nandodeomkar": pipeline("image-classification", 
+            model="nandodeomkar/autotrain-fracture-detection-using-google-vit-base-patch-16-54382127388"),
+        "Judy07": pipeline("object-detection", model="Judy07/bone-fracture-DETA")
     }
-    return models
 
-def draw_boxes(image, predictions):
+def draw_boxes(image, predictions, color):
     draw = ImageDraw.Draw(image)
     for pred in predictions:
         box = pred['box']
@@ -24,76 +49,89 @@ def draw_boxes(image, predictions):
         
         draw.rectangle(
             [(box['xmin'], box['ymin']), (box['xmax'], box['ymax'])],
-            outline="red",
-            width=3
+            outline=color,
+            width=2
         )
         
+        # Label background
         text_bbox = draw.textbbox((box['xmin'], box['ymin']), label)
-        draw.rectangle(text_bbox, fill="red")
+        draw.rectangle(text_bbox, fill=color)
         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")
+    st.title("🦴 Advanced Fracture Detection System")
     
     models = load_models()
     
-    uploaded_file = st.file_uploader("Upload X-ray image", type=['png', 'jpg', 'jpeg'])
+    with st.expander("⚙️ Settings", expanded=True):
+        conf_threshold = st.slider(
+            "Confidence threshold",
+            min_value=0.0,
+            max_value=1.0,
+            value=0.3,
+            step=0.01
+        )
     
-    conf_threshold = st.slider(
-        "Confidence threshold",
-        min_value=0.0,
-        max_value=1.0,
-        value=0.3,
-        step=0.01
+    uploaded_file = st.file_uploader(
+        "Upload X-ray image",
+        type=['png', 'jpg', 'jpeg'],
+        key="xray_upload"
     )
 
     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)
+        col1, col2 = st.columns([1, 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}")
-        
+            image = Image.open(uploaded_file)
+            max_size = (300, 300)
+            image.thumbnail(max_size, Image.Resampling.LANCZOS)
+            st.image(image, caption="Original X-ray", use_column_width=True)
+
         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")
+            tab1, tab2 = st.tabs(["📊 Classifications", "🔍 Detections"])
+            
+            with tab1:
+                for name in ["Heem2", "Nandodeomkar"]:
+                    with st.container():
+                        st.subheader(f"Model: {name}")
+                        with st.spinner("Analyzing..."):
+                            predictions = models[name](image)
+                            for pred in predictions:
+                                if pred['score'] >= conf_threshold:
+                                    score_color = "green" if pred['score'] > 0.7 else "orange"
+                                    st.markdown(f"""
+                                        <div style='padding: 10px; border-radius: 5px; background-color: #f0f2f6;'>
+                                            <span style='color: {score_color}; font-weight: bold;'>
+                                                {pred['score']:.1%}
+                                            </span> - {pred['label']}
+                                        </div>
+                                    """, unsafe_allow_html=True)
+
+            with tab2:
+                detection_colors = {"D3STRON": "#FF6B6B", "Judy07": "#4ECDC4"}
+                
+                for name, color in detection_colors.items():
+                    with st.container():
+                        st.subheader(f"Model: {name}")
+                        with st.spinner("Detecting..."):
+                            predictions = models[name](image)
+                            filtered_preds = [p for p in predictions if p['score'] >= conf_threshold]
+                            
+                            if filtered_preds:
+                                result_image = image.copy()
+                                result_image = draw_boxes(result_image, filtered_preds, color)
+                                st.image(result_image, use_column_width=True)
+                                
+                                for pred in filtered_preds:
+                                    st.markdown(f"""
+                                        <div style='padding: 8px; border-left: 4px solid {color}; 
+                                             margin: 5px 0; background-color: #f0f2f6;'>
+                                            {pred['label']}: {pred['score']:.1%}
+                                        </div>
+                                    """, unsafe_allow_html=True)
+                            else:
+                                st.info("No detections above threshold")
 
 if __name__ == "__main__":
     main()