Update app.py

app.py

@@ -82,12 +82,7 @@ def draw_bbox(image, bbox, label="Lesion", color="red", width=3):
     return image_copy
 
 def predict_skin_lesion(image):
-    
-    prompt = """Analyze this skin lesion image carefully. Look at the color, texture, shape, and borders. 
-    Provide your diagnosis and bounding box coordinates in this exact format: 
-    <diagnosis>condition_name</diagnosis> <bbox>[x1, y1, x2, y2]</bbox>
-    
-    Possible conditions: melanoma, melanocytic nevus, basal cell carcinoma, actinic keratosis, benign keratosis-like lesion, dermatofibroma, vascular lesion"""
+    prompt = "Analyze this skin lesion image. Respond in this exact format: <diagnosis>condition_name</diagnosis> <bbox>[x1, y1, x2, y2]</bbox>"
     
     conversation = [
         {
@@ -105,11 +100,11 @@ def predict_skin_lesion(image):
     with torch.no_grad():
         generated_ids = model.generate(
             **inputs,
-            max_new_tokens=150,
-            temperature=0.6,  # Lower temperature for more focused predictions
-            top_p=0.9,        
+            max_new_tokens=200,
+            temperature=0.3,
+            top_p=0.8,
+            # repetition_penalty=1.1,
             do_sample=True,
-            repetition_penalty=1.1,  # Reduce repetition
             pad_token_id=processor.tokenizer.eos_token_id
         )
     
@@ -122,23 +117,11 @@ def predict_skin_lesion(image):
     
     result_image = image.copy()
     if bbox:
-        result_image = draw_bbox(result_image, bbox, "AI Prediction", "red")
-    
-    # Better fallback handling
-    if diagnosis.lower() in dx_names:
-        full_diagnosis = dx_names[diagnosis.lower()]
-    elif any(dx in diagnosis.lower() for dx in dx_names.keys()):
-        for code, name in dx_names.items():
-            if code in diagnosis.lower():
-                full_diagnosis = name
-                break
-    else:
-        full_diagnosis = diagnosis.title()
+        result_image = draw_bbox(result_image, bbox, "Prediction", "red")
     
-    # Confidence indicator based on response quality
-    confidence = "High" if diagnosis and bbox else "Medium" if diagnosis else "Low"
+    full_diagnosis = dx_names.get(diagnosis.lower(), diagnosis)
     
-    return result_image, f"**Diagnosis:** {full_diagnosis}\n**Confidence:** {confidence}\n**Bbox:** {bbox}\n\n**Raw Response:** {response}"
+    return result_image, f"**Diagnosis:** {full_diagnosis}\n**Bbox:** {bbox}\n\n**Raw Response:** {response}"
 
 def load_random_and_predict():
     image, gt_diagnosis, gt_bbox = get_random_image()
@@ -159,33 +142,21 @@ def load_random_and_predict():
     
     return gt_image, result_image, prediction, ground_truth
 
-with gr.Blocks(title="DrDiag Skin Disease Diagnosis", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🔬 DrDiag: AI-Powered Skin Lesion Analysis")
-    gr.Markdown("Advanced skin disease diagnosis with bounding box localization using Qwen2-VL")
+with gr.Blocks(title="DrDiag Skin Disease Diagnosis") as demo:
+    gr.Markdown("# 🔬 DrDiag: Qwen2-VL Skin Disease Diagnosis with Spatial Awareness")
+    gr.Markdown("Analyze skin lesions with AI-powered diagnosis and bounding box detection")
     
     with gr.Row():
-        with gr.Column(scale=1):
-            input_image = gr.Image(type="pil", label="📤 Upload Skin Lesion Image", height=400)
-            
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Input Image")
             with gr.Row():
-                upload_btn = gr.Button("🔍 Analyze Image", variant="primary", size="lg")
-                random_btn = gr.Button("🎲 Random Sample", variant="secondary", size="lg")
-            
-            gr.Markdown("""
-            **📋 Detectable Conditions:**
-            - Melanoma
-            - Melanocytic Nevus  
-            - Basal Cell Carcinoma
-            - Actinic Keratosis
-            - Benign Keratosis-like Lesion
-            - Dermatofibroma
-            - Vascular Lesion
-            """)
+                upload_btn = gr.Button("Analyze Image", variant="primary")
+                random_btn = gr.Button("Random Sample", variant="secondary")
         
-        with gr.Column(scale=1):
-            output_image = gr.Image(type="pil", label="🎯 Analysis Result", height=400)
-            prediction_text = gr.Textbox(label="📊 AI Analysis", lines=8, show_copy_button=True)
-            ground_truth_text = gr.Textbox(label="✅ Ground Truth Comparison", lines=3, visible=False)
+        with gr.Column():
+            output_image = gr.Image(type="pil", label="Result with Bounding Box")
+            prediction_text = gr.Textbox(label="AI Prediction", lines=6)
+            ground_truth_text = gr.Textbox(label="Ground Truth", lines=4, visible=False)
     
     upload_btn.click(
         fn=predict_skin_lesion,