pre resnet model added

.gitattributes

@@ -34,3 +34,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 *.bmp filter=lfs diff=lfs merge=lfs -text
+*.t7 filter=lfs diff=lfs merge=lfs -text
+*.pth.tar filter=lfs diff=lfs merge=lfs -text
+*.txt filter=lfs diff=lfs merge=lfs -text

app.py

@@ -5,25 +5,85 @@ import torch
 from ultralytics import YOLO
 import numpy as np
 import os
+from torchvision import models, transforms
+import json
+
+# Load Pillow version
 from PIL import __version__ as PIL_VERSION
 print(f"Pillow version: {PIL_VERSION}")
 
 MODEL_PATH = "model/231220_detect_lr_0001_640_brightness.pt"
+SCENE_MODEL_PATH = "model/resnet50_places365.pth.tar"  # Updated path
+SCENE_LABELS_PATH = "model/categories_places365.txt"   # Updated path
 
 # Define the confidence threshold (used if not using the slider)
 # CONF_THRESHOLD = 0.5  # Optional: Remove if using the slider
 
-# Verify the model path
+# Verify the model paths
 if not os.path.exists(MODEL_PATH):
     raise FileNotFoundError(f"YOLO model not found at '{MODEL_PATH}'.")
+if not os.path.exists(SCENE_MODEL_PATH):
+    raise FileNotFoundError(f"Scene classification model not found at '{SCENE_MODEL_PATH}'.")
+if not os.path.exists(SCENE_LABELS_PATH):
+    raise FileNotFoundError(f"Scene classification labels not found at '{SCENE_LABELS_PATH}'.")
 
 # Load the YOLO model
 model = YOLO(MODEL_PATH)
 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)
+    checkpoint = torch.load(SCENE_MODEL_PATH, map_location=torch.device('cpu'))
+    state_dict = {str.replace(k, 'module.', ''): v for k, v in checkpoint['state_dict'].items()}
+    model.load_state_dict(state_dict)
+    model.eval()
+    return model
+
+scene_model = load_scene_classification_model()
+print("Scene classification model loaded.")
+
+# Load class labels
+with open(SCENE_LABELS_PATH) as class_file:
+    classes = class_file.read().splitlines()
+class_labels = [line.split(' ')[0][3:] for line in classes]  # Adjust parsing based on the file format
+
+# 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.5):
+    """
+    Classify the scene of the input image and check if it's a beach.
+    
+    Args:
+        input_image (PIL.Image): The uploaded image.
+        model (torch.nn.Module): The pre-trained scene classification model.
+        class_labels (list): List of class labels.
+        transform (torchvision.transforms): Image transformations.
+        threshold (float): Confidence threshold for beach classification.
+        
+    Returns:
+        bool: True if the image is classified as beach with confidence >= threshold, else False.
+        float: Confidence score for the beach classification.
+    """
+    image = transform(input_image).unsqueeze(0)  # Add batch dimension
+    with torch.no_grad():
+        outputs = model(image)
+        probabilities = torch.nn.functional.softmax(outputs, dim=1)
+        confidence, predicted = torch.max(probabilities, 1)
+        predicted_class = class_labels[predicted.item()]
+        is_beach = predicted_class.lower() in ['beach', 'seashore', 'shore', 'oceanfront'] and confidence.item() >= threshold
+    return is_beach, confidence.item()
+
 def detect_plastic_pellets(input_image, threshold=0.5):
     """
-    Perform plastic pellet detection using our customized model.
+    Perform plastic pellet detection using our customized model after verifying the scene.
     """
     if input_image is None:
         error_image = Image.new('RGB', (500, 100), color=(255, 0, 0))
@@ -36,7 +96,21 @@ def detect_plastic_pellets(input_image, threshold=0.5):
         return error_image
 
     try:
-        print("Starting detection with threshold:", threshold)
+        print("Starting scene classification...")
+        is_beach, scene_confidence = is_beach_scene(input_image, scene_model, class_labels, scene_transform, threshold=0.5)
+        
+        if not is_beach:
+            error_image = Image.new('RGB', (500, 100), color=(255, 165, 0))  # Orange color
+            draw = ImageDraw.Draw(error_image)
+            try:
+                font = ImageFont.truetype("arial.ttf", size=15)
+            except IOError:
+                font = ImageFont.load_default()
+            message = f"Image is not recognized as a beach (Confidence: {scene_confidence:.2f}). Please upload a beach image."
+            draw.text((10, 40), message, fill=(0, 0, 0), font=font)
+            return error_image
+
+        print("Scene classification passed. Starting detection...")
         input_image.thumbnail((1024, 1024), Image.LANCZOS)
         img = np.array(input_image.convert("RGB"))
 
@@ -135,4 +209,4 @@ def main():
     demo.launch()
 
 if __name__ == "__main__":
-        main()
+    main()

model/categories_places365.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2affba635eb657e7ca95f4e6cc69bd9fac29ef4c32aeb83cafdfcd06ec6a1ea6
+size 6833

model/resnet50_places365.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:46529c86902bd0cfb0ea562a30b2850c28d2620d96282b3db9c318e1d774f6c5
+size 97270159