sand class added to pre-checking model

app.py

@@ -19,8 +19,8 @@ print(f"Pillow version: {PIL_VERSION}")
 
 # Paths to models and labels
 MODEL_PATH = "model/231220_detect_lr_0001_640_brightness.pt"
-SCENE_MODEL_PATH = "model/resnet50_places365.pth.tar" 
-SCENE_LABELS_PATH = "model/categories_places365.txt"  
+SCENE_MODEL_PATH = "model/resnet50_places365.pth.tar"  # Updated path
+SCENE_LABELS_PATH = "model/categories_places365.txt"   # Updated path
 
 # Verify the model paths
 if not os.path.exists(MODEL_PATH):
@@ -37,13 +37,13 @@ print("YOLO model loaded.")
 # Load the scene classification model
 def load_scene_classification_model():
     # Load pre-trained ResNet50 model
-    model = models.resnet50(num_classes=365)
+    scene_model = models.resnet50(num_classes=365)
     checkpoint = torch.load(SCENE_MODEL_PATH, map_location=torch.device('cpu'))
     # Remove 'module.' prefix if present
     state_dict = {k.replace('module.', ''): v for k, v in checkpoint['state_dict'].items()}
-    model.load_state_dict(state_dict)
-    model.eval()
-    return model
+    scene_model.load_state_dict(state_dict)
+    scene_model.eval()
+    return scene_model
 
 scene_model = load_scene_classification_model()
 print("Scene classification model loaded.")
@@ -53,21 +53,23 @@ with open(SCENE_LABELS_PATH) as class_file:
     classes = class_file.read().splitlines()
 
 # Correct parsing of class labels
-class_labels = [line.split(' ', 1)[1].replace('_', ' ').lower() for line in classes]
+# Each line is in the format '/a/beach 48', so we extract 'beach'
+class_labels = [line.split(' ')[0][3:].lower() for line in classes]
 
 # Debug: Print some class labels to verify parsing
 print("Sample Class Labels:")
 for idx in range(10):
     print(f"{idx}: {class_labels[idx]}")
 
-# Define beach-related keywords for flexible matching
-beach_keywords = [
-    'beach', 'seashore', 'shore', 'oceanfront', 'sandy', 'seaside',
-    'coast', 'island', 'rocky', 'tropical', 'surf', 'resort',
-    'sunset', 'sunrise', 'sand'
-]
+# Define image transformations for scene classification
+scene_transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],  # ImageNet means
+                         std=[0.229, 0.224, 0.225])   # ImageNet stds
+])
 
-def is_beach_scene(input_image, model, class_labels, transform, threshold=0.1):
+def is_beach_scene(input_image, model, class_labels, transform, threshold=0.2):
     """
     Classify the scene of the input image and check if it's a beach.
     
@@ -88,10 +90,13 @@ def is_beach_scene(input_image, model, class_labels, transform, threshold=0.1):
         probabilities = torch.nn.functional.softmax(outputs, dim=1)
         confidence, predicted = torch.max(probabilities, 1)
         predicted_class = class_labels[predicted.item()]
+        predicted_class_lower = predicted_class.lower()
+        
+        # Check if 'beach' or 'sand' is in the predicted class and exclude 'desert'
+        is_beach = (('beach' in predicted_class_lower or 'sand' in predicted_class_lower) and 
+                   ('desert' not in predicted_class_lower) and 
+                   confidence.item() >= threshold)
         
-        # Flexible matching using regex for whole words
-        is_beach = any(re.search(r'\b' + re.escape(keyword) + r'\b', predicted_class) for keyword in beach_keywords) and confidence.item() >= threshold
-
     # Log the classification result
     logging.info(f"Predicted Class: {predicted_class}, Confidence: {confidence.item():.4f}, Is Beach: {is_beach}")
     
@@ -101,7 +106,7 @@ def is_beach_scene(input_image, model, class_labels, transform, threshold=0.1):
     
     return is_beach, confidence.item()
 
-def detect_plastic_pellets(input_image, scene_threshold=0.1, detection_threshold=0.5):
+def detect_plastic_pellets(input_image, threshold=0.5):
     """
     Perform plastic pellet detection using our customized model after verifying the scene.
     """
@@ -117,30 +122,24 @@ def detect_plastic_pellets(input_image, scene_threshold=0.1, detection_threshold
         return error_image
 
     try:
-        logging.info(f"Starting scene classification with threshold: {scene_threshold}")
-        print(f"Starting scene classification with threshold: {scene_threshold}")
-        is_beach, scene_confidence = is_beach_scene(
-            input_image, 
-            scene_model, 
-            class_labels, 
-            scene_transform, 
-            threshold=scene_threshold
-        )
-
+        print("Starting scene classification...")
+        logging.info("Starting scene classification...")
+        is_beach, scene_confidence = is_beach_scene(input_image, scene_model, class_labels, scene_transform, threshold=0.2)
+        
         if not is_beach:
-            logging.warning("Image not recognized as beach.")
+            logging.warning("Image not recognized as a beach.")
             error_image = Image.new('RGB', (500, 150), color=(255, 165, 0))  # Increased height for more text
             draw = ImageDraw.Draw(error_image)
             try:
                 font = ImageFont.truetype("arial.ttf", size=15)
             except IOError:
                 font = ImageFont.load_default()
-            message = f"Image not recognized as a beach.\nConfidence: {scene_confidence:.2f}"
+            message = f"Image is not recognized as a beach.\nConfidence: {scene_confidence:.2f}"
             draw.text((10, 40), message, fill=(0, 0, 0), font=font)
             return error_image
 
-        logging.info("Scene classification passed. Starting detection...")
         print("Scene classification passed. Starting detection...")
+        logging.info("Scene classification passed. Starting detection...")
         input_image.thumbnail((1024, 1024), Image.LANCZOS)
         img = np.array(input_image.convert("RGB"))
 
@@ -156,7 +155,7 @@ def detect_plastic_pellets(input_image, scene_threshold=0.1, detection_threshold
         for result in results:
             for box in result.boxes:
                 confidence = box.conf[0].item()
-                if confidence < detection_threshold:
+                if confidence < threshold:
                     continue  # Skip detections below the threshold
 
                 x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
@@ -175,13 +174,14 @@ def detect_plastic_pellets(input_image, scene_threshold=0.1, detection_threshold
 
         if detection_made:
             logging.info("Plastic pellets detected.")
+            print("Plastic pellets detected.")
         else:
             logging.info("No plastic pellets detected.")
             draw.text((10, 10), "No plastic pellets detected.", fill=(255, 0, 0), font=font)
             return input_image  
 
-        logging.info("Detection completed.")
         print("Detection completed.")
+        logging.info("Detection completed.")
         return input_image
 
     except Exception as e:
@@ -211,23 +211,13 @@ def main():
                 examples = ['images/image1.bmp', 'images/image2.bmp', 'images/image3.bmp']
                 gr.Examples(examples=examples, inputs=input_image, label="Or choose one of these images")
                 
-                # Add a slider for Scene Classification Threshold
-                scene_threshold = gr.Slider(
-                    minimum=0.0,
-                    maximum=1.0,
-                    value=0.1,  # Default value set to 0.1
-                    step=0.05,
-                    label="Scene Classification Threshold",
-                    info="Adjust the confidence threshold for scene classification (pre-check)."
-                )
-                
-                # Add a slider for Detection Confidence Threshold
-                detection_threshold = gr.Slider(
+                # Add a slider for confidence threshold
+                confidence_threshold = gr.Slider(
                     minimum=0.0,
                     maximum=1.0,
-                    value=0.5,  # Default value remains at 0.5
+                    value=0.5,
                     step=0.05,
-                    label="Detection Confidence Threshold",
+                    label="Confidence Threshold",
                     info="Adjust the confidence threshold for displaying detections."
                 )
                 
@@ -245,7 +235,7 @@ def main():
 
         submit_button.click(
             fn=detect_plastic_pellets,
-            inputs=[input_image, scene_threshold, detection_threshold],
+            inputs=[input_image, confidence_threshold],
             outputs=output_image,
             api_name="detect",
             show_progress=True