Update app.py

app.py

@@ -35,7 +35,7 @@ except Exception as e:
     print(f"Salesforce connection failed: {str(e)}")
     raise
 
-# Load pre-trained model and processor (replace with a construction-specific model if available)
+# Load pre-trained model and processor
 model_name = "google/vit-base-patch16-224"  # Placeholder; replace with "nasreshsuguru/construction-milestone-detector" if trained
 processor = ViTImageProcessor.from_pretrained(model_name)
 model = ViTForImageClassification.from_pretrained(model_name)
@@ -48,18 +48,18 @@ def process_image(image, project_name):
     try:
         # Validate inputs
         if image is None:
-            return "Error: Please upload an image to proceed.", "Pending", "", 0.0, 0.0, ""
+            return "Error: Please upload an image to proceed.", "Pending", "", ""
         if not project_name:
-            return "Error: Please enter a project name to proceed.", "Pending", "", 0.0, 0.0, ""
+            return "Error: Please enter a project name to proceed.", "Pending", "", ""
         if not re.match(r'^[a-zA-Z0-9\s-]+$', project_name):
-            return "Error: Project name must be alphanumeric (letters, numbers, spaces, or hyphens).", "Pending", "", 0.0, 0.0, ""
+            return "Error: Project name must be alphanumeric (letters, numbers, spaces, or hyphens).", "Pending", "", ""
 
         # Validate image size and type
         image_size_mb = os.path.getsize(image) / (1024 * 1024)
         if image_size_mb > 20:
-            return "Error: Image size exceeds 20MB.", "Failure", "", 0.0, 0.0, ""
+            return "Error: Image size exceeds 20MB.", "Failure", "", ""
         if not image.lower().endswith(('.jpg', '.jpeg', '.png')):
-            return "Error: Only JPG/PNG images are supported.", "Failure", "", 0.0, 0.0, ""
+            return "Error: Only JPG/PNG images are supported.", "Failure", "", ""
 
         # Preprocess image
         img = Image.open(image).convert("RGB")
@@ -80,21 +80,18 @@ def process_image(image, project_name):
         
         # Map indices to labels
         predicted_idx = top_indices[0]
-        confidence = top_probs[0]
         milestone = model.config.id2label.get(predicted_idx, "Unknown Milestone")
-        completion_percent = min(max(int(confidence * 100), 0), 100)
         
         # Format top predictions
-        prediction_details = "\n".join([f"{model.config.id2label.get(idx, 'Unknown')}: {prob:.2f}" for idx, prob in zip(top_indices, top_probs)])
+        prediction_details = "\n".join([f"{model.config.id2label.get(idx, 'Unknown Milestone')}: {prob:.2f}" for idx, prob in zip(top_indices, top_probs)])
 
         # Update Salesforce record
         record = {
             "Name": project_name,
             "Current_Milestone__c": milestone,
-            "Completion_Percentage__c": completion_percent,
             "Last_Updated_On__c": datetime.now().isoformat(),
             "Upload_Status__c": "Success",
-            "Comments__c": f"AI Confidence: {confidence:.2f}",
+            "Comments__c": f"AI Prediction: {milestone}",
             "Version__c": 1
         }
 
@@ -110,19 +107,17 @@ def process_image(image, project_name):
             else:
                 sf.Construction_Project__c.create(record)
         except Exception as e:
-            return f"Error: Failed to update Salesforce - {str(e)}", "Failure", "", 0.0, 0.0, prediction_details
+            return f"Error: Failed to update Salesforce - {str(e)}", "Failure", "", prediction_details
 
         return (
-            f"Success: Milestone: {milestone}, Completion: {completion_percent}%",
+            f"Success: Milestone: {milestone}",
             "Success",
             milestone,
-            completion_percent,
-            confidence,
             prediction_details
         )
 
     except Exception as e:
-        return f"Error: {str(e)}", "Failure", "", 0.0, 0.0, ""
+        return f"Error: {str(e)}", "Failure", "", ""
 
 # Gradio interface
 with gr.Blocks(css=".gradio-container {background-color: #f0f4f8; font-family: Arial;} .title {color: #2c3e50; font-size: 24px; text-align: center;}") as demo:
@@ -133,8 +128,6 @@ with gr.Blocks(css=".gradio-container {background-color: #f0f4f8; font-family: A
     output_text = gr.Textbox(label="Result")
     upload_status = gr.Textbox(label="Upload Status")
     milestone = gr.Textbox(label="Detected Milestone")
-    completion_percent = gr.Slider(0, 100, label="Completion Percentage (%)", interactive=False)
-    confidence_score = gr.Slider(0, 1, label="Confidence Score", interactive=False)
     prediction_details = gr.Textbox(label="Top Predictions")
     progress = gr.Slider(0, 100, label="Processing Progress", interactive=False, value=0)
 
@@ -150,11 +143,10 @@ with gr.Blocks(css=".gradio-container {background-color: #f0f4f8; font-family: A
     ).then(
         fn=process_image,
         inputs=[image_input, project_name],
-        outputs=[output_text, upload_status, milestone, completion_percent, confidence_score, prediction_details]
+        outputs=[output_text, upload_status, milestone, prediction_details]
     ).then(
         fn=complete_progress,
         outputs=progress
     )
 
-if __name__ == "__main__":
-    demo.launch()
+if __name__ == "__main__":