Update app.py

app.py

@@ -225,26 +225,30 @@ def classify(UserInput, Image, Textbox2, Textbox3):
         
             if max_label_index is not None:
                 max_label = labels[max_label_index].split(' ', 1)[1]
-                print(f'Maximum Prediction: {max_label} with a value of {round(max_prediction_value, 2)}')
+                rounded_prediction = round(max_prediction_value, 2)
+                print(f'Maximum Prediction: {max_label} with a value of {rounded_prediction}')
     
                 time.sleep(1)
-                print("\nWays to dispose of this waste: " + max_label)
-                messages.append({"role": "user", "content": content + " " + max_label})
-    
-                headers = {
-                    "Content-Type": "application/json",
-                    "Authorization": f"Bearer {auth}"
-                }
+                if rounded_prediction > 0.5:
+                    print("\nWays to dispose of this waste: " + max_label)
+                    messages.append({"role": "user", "content": content + " " + max_label})
         
-                response = requests.post(host, headers=headers, json={
-                    "messages": messages,
-                    "model": model_llm
-                }).json()
-                
-                reply = response["choices"][0]["message"]["content"]
-                messages.append({"role": "assistant", "content": reply})
-    
-                output.append({"Mode": "Image", "type": max_label, "prediction_value": rounded_value, "content": reply})
+                    headers = {
+                        "Content-Type": "application/json",
+                        "Authorization": f"Bearer {auth}"
+                    }
+            
+                    response = requests.post(host, headers=headers, json={
+                        "messages": messages,
+                        "model": model_llm
+                    }).json()
+                    
+                    reply = response["choices"][0]["message"]["content"]
+                    messages.append({"role": "assistant", "content": reply})
+                    
+                    output.append({"Mode": "Image", "type": max_label, "prediction_value": rounded_value, "content": reply})
+                elif rounded_prediction < 0.5:
+                    output.append({"Mode": "Image", "type": "Not predictable", "prediction_value": rounded_value, "content": "Seems like the prediction rate is too low due to that won't be able to predict the type of material. Try again with a cropped image or different one."})
             
             return output