resnet deleted

app.py

@@ -5,113 +5,27 @@ import torch
 from ultralytics import YOLO
 import numpy as np
 import os
-from torchvision import models, transforms
-import re
-import logging
-
-# Configure logging
-logging.basicConfig(filename='app.log', level=logging.INFO,
-                    format='%(asctime)s:%(levelname)s:%(message)s')
-
-# Load Pillow version
 from PIL import __version__ as PIL_VERSION
 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"  # Updated path
-SCENE_LABELS_PATH = "model/categories_places365.txt"   # Updated path
 
-# Verify the model paths
+# Define the confidence threshold (used if not using the slider)
+# CONF_THRESHOLD = 0.5  # Optional: Remove if using the slider
+
+# Verify the model path
 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
-    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()}
-    scene_model.load_state_dict(state_dict)
-    scene_model.eval()
-    return scene_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()
-
-# Correct parsing of class labels
-# 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 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.2):
-    """
-    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()]
-        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)
-        
-    # Log the classification result
-    logging.info(f"Predicted Class: {predicted_class}, Confidence: {confidence.item():.4f}, Is Beach: {is_beach}")
-    
-    # Debug: Print predicted class and confidence
-    print(f"Predicted Class: {predicted_class}, Confidence: {confidence.item():.4f}")
-    print(f"Is Beach: {is_beach}")
-    
-    return is_beach, confidence.item()
-
 def detect_plastic_pellets(input_image, threshold=0.5):
     """
-    Perform plastic pellet detection using our customized model after verifying the scene.
+    Perform plastic pellet detection using our customized model.
     """
     if input_image is None:
-        logging.warning("No image uploaded.")
         error_image = Image.new('RGB', (500, 100), color=(255, 0, 0))
         draw = ImageDraw.Draw(error_image)
         try:
@@ -122,24 +36,7 @@ def detect_plastic_pellets(input_image, threshold=0.5):
         return error_image
 
     try:
-        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 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 is not recognized as a beach.\nConfidence: {scene_confidence:.2f}"
-            draw.text((10, 40), message, fill=(0, 0, 0), font=font)
-            return error_image
-
-        print("Scene classification passed. Starting detection...")
-        logging.info("Scene classification passed. Starting detection...")
+        print("Starting detection with threshold:", threshold)
         input_image.thumbnail((1024, 1024), Image.LANCZOS)
         img = np.array(input_image.convert("RGB"))
 
@@ -172,20 +69,14 @@ def detect_plastic_pellets(input_image, threshold=0.5):
 
                 detection_made = True
 
-        if detection_made:
-            logging.info("Plastic pellets detected.")
-            print("Plastic pellets detected.")
-        else:
-            logging.info("No plastic pellets detected.")
+        if not detection_made:
             draw.text((10, 10), "No plastic pellets detected.", fill=(255, 0, 0), font=font)
             return input_image  
 
         print("Detection completed.")
-        logging.info("Detection completed.")
         return input_image
 
     except Exception as e:
-        logging.error(f"Detection error: {str(e)}")
         print(f"Detection error: {str(e)}")
         error_image = Image.new('RGB', (500, 100), color=(255, 0, 0))
         draw = ImageDraw.Draw(error_image)
@@ -244,4 +135,4 @@ def main():
     demo.launch()
 
 if __name__ == "__main__":
-    main()
+        main()

model/categories_places365.txt

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

model/resnet50_places365.pth.tar

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

requirements.txt

@@ -1,6 +1,3 @@
-gradio>=3.38.0
-torch>=2.0.0
-torchvision>=0.15.1
-ultralytics>=8.0.0
-pillow>=10.0.0
-numpy>=1.23.0
+version https://git-lfs.github.com/spec/v1
+oid sha256:0689b23c5d7d1c089c59d97ac59bee19bec098c7857c300e9df9815cc1840d63
+size 96