update

.gitignore

@@ -1 +1,2 @@
-app_backup.py
+app_backup.py
+app_backup_original.py

app.py

@@ -69,7 +69,8 @@ st.markdown("""
 @st.cache_resource
 def load_model():
     REVISION = 'refs/pr/6'
-    MODEL_NAME = "Anonymous-AC/AD-KD"
+    # MODEL_NAME = "RioJune/AD-KD-MICCAI25"
+    MODEL_NAME = '/u/home/lj0/Checkpoints/AD-KD-MICCAI25'
     DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     config_model = AutoConfig.from_pretrained("microsoft/Florence-2-base-ft", trust_remote_code=True)
@@ -155,10 +156,12 @@ with col2:
 
     st.subheader("⚠️ Warning:")
     st.write("""
-        - **🚫 Please avoid uploading non-frontal chest X-ray images**. Our model has been specifically trained on **frontal chest X-ray images**.
-        - This demo is intended for **🔬 research purposes only** and should **❌ not be used for medical diagnoses**.
-        - The model’s responses may contain **🤖 hallucinations or incorrect information**. Always consult a **👨‍⚕️ medical professional** for accurate diagnosis and advice.
-    """)
+    - **🚫 Please avoid uploading non-frontal chest X-ray images.** Our model has been specifically trained on **frontal chest X-ray images** only.
+    - This demo is intended for **🔬 research purposes only** and should **❌ not be used for medical diagnoses**.
+    - The model’s responses may contain **<span style='color:#dc3545; font-weight:bold;'>🤖 hallucinations or incorrect information</span>**. 
+    - Always consult a **<span style='color:#dc3545; font-weight:bold;'>👨‍⚕️ medical professional</span>** for accurate diagnosis and advice.
+""", unsafe_allow_html=True)
+
 
     st.markdown("</div>", unsafe_allow_html=True)
 
@@ -182,10 +185,45 @@ if st.button("Run Inference on Example", key="example"):
         inputs = processor(text=[prompt], images=[np_image], return_tensors="pt", padding=True).to(DEVICE)
 
         with st.spinner("Processing... ⏳"):
-            # Generate the result
-            generated_ids = model.generate(input_ids=inputs["input_ids"], pixel_values=inputs["pixel_values"], max_new_tokens=1024, num_beams=3)
+            outputs = model.generate(
+                input_ids=inputs["input_ids"],
+                pixel_values=inputs["pixel_values"],
+                max_new_tokens=1024,
+                num_beams=3,
+                output_scores=True,  # Make sure we get the scores/logits
+                return_dict_in_generate=True  # Ensures you get both sequences and scores in the output
+                )
+            
+
+            # Ensure transition_scores is properly extracted
+            transition_scores = model.compute_transition_scores(
+                outputs.sequences, outputs.scores, outputs.beam_indices, normalize_logits=False
+            )
+
+            # Get the generated token IDs (ignoring the input tokens part)
+            generated_ids = outputs.sequences
             generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
 
+            # Get input length
+            input_length = inputs.input_ids.shape[1]
+            generated_tokens = outputs.sequences
+
+            # Calculate output length (number of generated tokens)
+            output_length = np.sum(transition_scores.cpu().numpy() < 0, axis=1)
+
+            # Get length penalty
+            length_penalty = model.generation_config.length_penalty
+
+            # Calculate total score for the generated sentence
+            reconstructed_scores = transition_scores.cpu().sum(axis=1) / (output_length**length_penalty)
+
+            # Convert log-probability to probability (0-1 range)
+            probabilities = np.exp(reconstructed_scores.cpu().numpy())
+
+            # Streamlit UI to display the result
+            st.markdown(f"**🎯 Probability of the Results:** <span style='color:#28a745; font-size:24px; font-weight:bold;'>{probabilities[0] * 100:.2f}%</span>", unsafe_allow_html=True)
+                
+
             predictions = processor.post_process_generation(generated_text, task="<CAPTION_TO_PHRASE_GROUNDING>", image_size=np_image.shape[:2])
 
             detection = sv.Detections.from_lmm(sv.LMM.FLORENCE_2, predictions, resolution_wh=np_image.shape[:2])
@@ -221,18 +259,6 @@ with col2:
 
 uploaded_file = st.file_uploader("Upload an Image", type=["png", "jpg", "jpeg"])
 
-# if uploaded_file:
-#     image = Image.open(uploaded_file).convert("RGB")
-#     image = apply_transform(image)  # Ensure the uploaded image is transformed correctly
-#     st.image(image, caption="Uploaded Image", width=400)
-
-#     # Let user select dataset and disease dynamically
-#     disease_choice = disease_choice if disease_choice else example_diseases[0]
-
-# # Get Definition Priority: Dataset -> User Input
-# definition = vindr_definitions.get(disease_choice, padchest_definitions.get(disease_choice, ""))
-# if not definition:
-#     definition = st.text_input("Enter Definition Manually 📝", value="")
 
 col1, col2 = st.columns([1, 2])
 
@@ -264,10 +290,11 @@ with col2:
 
     st.subheader("⚠️ Warning:")
     st.write("""
-        - **🚫 Please avoid uploading non-frontal chest X-ray images**. Our model has been specifically trained on **frontal chest X-ray images**.
-        - This demo is intended for **🔬 research purposes only** and should **❌ not be used for medical diagnoses**.
-        - The model’s responses may contain **🤖 hallucinations or incorrect information**. Always consult a **👨‍⚕️ medical professional** for accurate diagnosis and advice.
-    """)
+    - **🚫 Please avoid uploading non-frontal chest X-ray images.** Our model has been specifically trained on **frontal chest X-ray images** only.
+    - This demo is intended for **🔬 research purposes only** and should **❌ not be used for medical diagnoses**.
+    - The model’s responses may contain **<span style='color:#dc3545; font-weight:bold;'>🤖 hallucinations or incorrect information</span>**. 
+    - Always consult a **<span style='color:#dc3545; font-weight:bold;'>👨‍⚕️ medical professional</span>** for accurate diagnosis and advice.
+""", unsafe_allow_html=True)
 
 # Run inference after upload
 if st.button("Run Inference 🏃‍♂️"):
@@ -285,9 +312,52 @@ if st.button("Run Inference 🏃‍♂️"):
         inputs = processor(text=[prompt], images=[np_image], return_tensors="pt", padding=True).to(DEVICE)
 
         with st.spinner("Processing... ⏳"):
-            generated_ids = model.generate(input_ids=inputs["input_ids"], pixel_values=inputs["pixel_values"], max_new_tokens=1024, num_beams=3)
+            # generated_ids = model.generate(input_ids=inputs["input_ids"], pixel_values=inputs["pixel_values"], max_new_tokens=1024, num_beams=3)
+            # generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+            
+            outputs = model.generate(
+                input_ids=inputs["input_ids"],
+                pixel_values=inputs["pixel_values"],
+                max_new_tokens=1024,
+                num_beams=3,
+                output_scores=True,  # Make sure we get the scores/logits
+                return_dict_in_generate=True  # Ensures you get both sequences and scores in the output
+                )
+
+            transition_scores = model.compute_transition_scores(
+                outputs.sequences, outputs.scores, outputs.beam_indices, normalize_logits=False
+            )
+
+            # Get the generated token IDs (ignoring the input tokens part)
+            generated_ids = outputs.sequences
             generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
 
+            # Get input length
+            input_length = inputs.input_ids.shape[1]
+
+            # Extract generated tokens (ignoring the input tokens)
+            # generated_tokens = outputs.sequences[:, input_length:]
+            generated_tokens = outputs.sequences
+
+            # Calculate output length (number of generated tokens)
+            output_length = np.sum(transition_scores.cpu().numpy() < 0, axis=1)
+
+            # Get length penalty
+            length_penalty = model.generation_config.length_penalty
+
+            # Calculate total score for the generated sentence
+            reconstructed_scores = transition_scores.cpu().sum(axis=1) / (output_length**length_penalty)
+
+            # Convert log-probability to probability (0-1 range)
+            probabilities = np.exp(reconstructed_scores.cpu().numpy())
+
+            # Streamlit UI to display the result
+            
+            # st.write(f"**Probability of the Results (0-1):** {probabilities[0]:.4f}")
+            st.markdown(f"**🎯 Probability of the Results:** <span style='color:green; font-size:24px; font-weight:bold;'>{probabilities[0] * 100:.2f}%</span>", unsafe_allow_html=True)
+
+
+
             predictions = processor.post_process_generation(generated_text, task="<CAPTION_TO_PHRASE_GROUNDING>", image_size=np_image.shape[:2])
 
             detection = sv.Detections.from_lmm(sv.LMM.FLORENCE_2, predictions, resolution_wh=np_image.shape[:2])
@@ -311,3 +381,6 @@ if st.button("Run Inference 🏃‍♂️"):
 
             # Display the generated text
             st.write("**Generated Text:**", generated_text)
+
+
+