Update app.py

app.py

@@ -33,25 +33,10 @@ except Exception as e:
     logging.error(f"Failed to connect to Salesforce: {str(e)}")
     raise Exception(f"Failed to connect to Salesforce: {str(e)}")
 
-# Load the Faster R-CNN pretrained model (replace with your fine-tuned weights if any)
+# Load the Faster R-CNN pretrained model
 model = models.detection.fasterrcnn_resnet50_fpn(weights="FasterRCNN_ResNet50_FPN_Weights.COCO_V1")
 model.eval()
 
-# COCO categories list for mapping labels (standard)
-COCO_INSTANCE_CATEGORY_NAMES = [
-    '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
-    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter',
-    'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra',
-    'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
-    'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard',
-    'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon',
-    'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
-    'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet',
-    'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven',
-    'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
-    'hair drier', 'toothbrush'
-]
-
 # Image transformation for the model input
 transform = transforms.Compose([
     transforms.ToTensor(),
@@ -66,17 +51,10 @@ def get_severity(score):
     else:
         return "Minor"
 
-# Temporary mapping from COCO labels to defect types (replace with your own)
-COCO_TO_DEFECT_MAPPING = {
-    'car': 'Crack',
-    'person': 'Rust',
-    'bicycle': 'Deformation',
-    'truck': 'Corrosion',
-    'boat': 'Spalling',
-}
-
-def map_defect_type(coco_label):
-    return COCO_TO_DEFECT_MAPPING.get(coco_label, "Crack")
+# Simple defect type mapping (no COCO labels)
+# Here, just return a generic "Defect" label or you can customize per your need
+def map_defect_type():
+    return "Defect"
 
 # Upload annotated image to Salesforce as ContentVersion record
 def upload_image_to_salesforce(image, filename="detected_image.jpg", record_id=None):
@@ -113,7 +91,7 @@ def detect_defects(image):
         result_image = image.copy()
         draw = ImageDraw.Draw(result_image)
 
-        # Optional: Use a truetype font for nicer text, fallback if not available
+        # Use default font
         try:
             font = ImageFont.truetype("arial.ttf", 18)
         except:
@@ -126,10 +104,8 @@ def detect_defects(image):
                 continue
 
             box = predictions[0]['boxes'][i].tolist()
-            label_idx = predictions[0]['labels'][i].item()
-            coco_label = COCO_INSTANCE_CATEGORY_NAMES[label_idx]
 
-            defect_type = map_defect_type(coco_label)
+            defect_type = map_defect_type()
             severity = get_severity(score)
 
             # Append defect info to output list
@@ -137,7 +113,6 @@ def detect_defects(image):
                 "type": defect_type,
                 "confidence": round(score, 2),
                 "severity": severity,
-                "coco_label": coco_label
             })
 
             # Draw rectangle and label