bplusplus 1.2.1__py3-none-any.whl → 1.2.3__py3-none-any.whl

bplusplus/resnet/test.py

@@ -1,473 +0,0 @@
-# pip install ultralytics torchvision pillow numpy scikit-learn tabulate tqdm
-# python3 tests/two-stage(yolo-resnet).py --data ' --yolo_weights --resnet_weights --use_resnet50
-
-import os
-import cv2
-import torch
-from ultralytics import YOLO
-from torchvision import transforms
-from PIL import Image
-import numpy as np
-from torchvision.models import resnet152, resnet50
-import torch.nn as nn
-from sklearn.metrics import classification_report, accuracy_score
-import time
-from collections import defaultdict
-from tabulate import tabulate
-from tqdm import tqdm
-import csv
-import requests
-import sys
-
-def test_resnet(data_path, yolo_weights, resnet_weights, model="resnet152", species_names=None, output_dir="output"):
-    """
-    Run the two-stage detection and classification test
-    
-    Args:
-        data_path (str): Path to the test directory
-        yolo_weights (str): Path to the YOLO model file
-        resnet_weights (str): Path to the ResNet model file
-        model (str): Model type, either "resnet50" or "resnet152"
-        species_names (list): List of species names
-        output_dir (str): Directory to save output CSV files
-    """    
-    use_resnet50 = model == "resnet50"
-    classifier = TestTwoStage(yolo_weights, resnet_weights, use_resnet50=use_resnet50, 
-                             species_names=species_names, output_dir=output_dir)
-    classifier.run(data_path)
-
-class TestTwoStage:
-    def __init__(self, yolo_model_path, resnet_model_path, num_classes=9, use_resnet50=False, species_names="", output_dir="output"):
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        print(f"Using device: {self.device}")
-        
-        self.output_dir = output_dir
-        os.makedirs(self.output_dir, exist_ok=True)
-        
-        self.yolo_model = YOLO(yolo_model_path)
-        self.classification_model = resnet50(pretrained=False) if use_resnet50 else resnet152(pretrained=False)
-    
-        self.classification_model.fc = nn.Sequential(
-            nn.Dropout(0.4),  # Using dropout probability of 0.4 as in training
-            nn.Linear(self.classification_model.fc.in_features, num_classes)
-        )
-        
-        state_dict = torch.load(resnet_model_path, map_location=self.device)
-        self.classification_model.load_state_dict(state_dict)
-        self.classification_model.to(self.device)
-        self.classification_model.eval()
-
-        self.classification_transform = transforms.Compose([
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-        ])
-        self.species_names = species_names
-        
-    def get_frames(self, test_dir):
-        image_dir = os.path.join(test_dir, "images")
-        label_dir = os.path.join(test_dir, "labels")
-        
-        predicted_frames = []
-        true_frames = []
-        image_names = []
-
-        start_time = time.time()  # Start timing
-
-        for image_name in tqdm(os.listdir(image_dir), desc="Processing Images", unit="image"):
-            image_names.append(image_name)
-            image_path = os.path.join(image_dir, image_name)
-            label_path = os.path.join(label_dir, image_name.replace('.jpg', '.txt'))
-
-            frame = cv2.imread(image_path)
-            # Suppress print statements from YOLO model
-            with torch.no_grad():
-                results = self.yolo_model(frame, conf=0.3, iou=0.5, verbose=False)
-
-            detections = results[0].boxes
-            predicted_frame = []
-
-            if detections:
-                for box in detections:
-                    xyxy = box.xyxy.cpu().numpy().flatten()
-                    x1, y1, x2, y2 = xyxy[:4]
-                    width = x2 - x1
-                    height = y2 - y1
-                    x_center = x1 + width / 2
-                    y_center = y1 + height / 2
-
-                    insect_crop = frame[int(y1):int(y2), int(x1):int(x2)]
-                    insect_crop_rgb = cv2.cvtColor(insect_crop, cv2.COLOR_BGR2RGB)
-                    pil_img = Image.fromarray(insect_crop_rgb)
-                    input_tensor = self.classification_transform(pil_img).unsqueeze(0).to(self.device)
-
-                    with torch.no_grad():
-                        outputs = self.classification_model(input_tensor)
-
-                    # Directly use the model output without any mapping
-                    predicted_class_idx = outputs.argmax(dim=1).item()
-                    
-                    img_height, img_width, _ = frame.shape
-                    x_center_norm = x_center / img_width
-                    y_center_norm = y_center / img_height
-                    width_norm = width / img_width
-                    height_norm = height / img_height
-                    predicted_frame.append([predicted_class_idx, x_center_norm, y_center_norm, width_norm, height_norm])
-
-            predicted_frames.append(predicted_frame if predicted_frame else [])
-
-            true_frame = []
-            if os.path.exists(label_path) and os.path.getsize(label_path) > 0:
-                with open(label_path, 'r') as f:
-                    for line in f:
-                        label_line = line.strip().split()
-                        true_frame.append([int(label_line[0]), *map(np.float32, label_line[1:])])
-
-            true_frames.append(true_frame if true_frame else [])
-
-        end_time = time.time()  # End timing
-
-        model_type = "resnet50" if isinstance(self.classification_model, type(resnet50())) else "resnet152"
-        output_file = os.path.join(self.output_dir, f"results_{model_type}_{time.strftime('%Y%m%d_%H%M%S')}.csv")
-        
-        with open(output_file, "w") as f:
-            writer = csv.writer(f)
-            writer.writerow(["Image Name", "True", "Predicted"])
-            for image_name, true_frame, predicted_frame in zip(image_names, true_frames, predicted_frames):
-                writer.writerow([image_name, true_frame, predicted_frame])
-        
-        print(f"Results saved to {output_file}")
-        return predicted_frames, true_frames, end_time - start_time
-    
-    def get_taxonomic_info(self, species_list):
-        """
-        Retrieves taxonomic information for a list of species from GBIF API.
-        Creates a hierarchical taxonomy dictionary with order, family, and species relationships.
-        """
-        taxonomy = {1: [], 2: {}, 3: {}}
-        species_to_family = {}
-        family_to_order = {}
-        
-        print(f"Building taxonomy from GBIF for {len(species_list)} species")
-        
-        print("\nTaxonomy Results:")
-        print("-" * 80)
-        print(f"{'Species':<30} {'Order':<20} {'Family':<20} {'Status'}")
-        print("-" * 80)
-        
-        for species_name in species_list:
-            url = f"https://api.gbif.org/v1/species/match?name={species_name}&verbose=true"
-            try:
-                response = requests.get(url)
-                data = response.json()
-                
-                if data.get('status') == 'ACCEPTED' or data.get('status') == 'SYNONYM':
-                    family = data.get('family')
-                    order = data.get('order')
-                    
-                    if family and order:
-                        status = "OK"
-                        
-                        print(f"{species_name:<30} {order:<20} {family:<20} {status}")
-                        
-                        species_to_family[species_name] = family
-                        family_to_order[family] = order
-                        
-                        if order not in taxonomy[1]:
-                            taxonomy[1].append(order)
-                        
-                        taxonomy[2][family] = order
-                        taxonomy[3][species_name] = family
-                    else:
-                        error_msg = f"Species '{species_name}' found in GBIF but family and order not found, could be spelling error in species, check GBIF"
-                        print(f"{species_name:<30} {'Not found':<20} {'Not found':<20} ERROR")
-                        print(f"Error: {error_msg}")
-                        sys.exit(1)  # Stop the script
-                else:
-                    error_msg = f"Species '{species_name}' not found in GBIF, could be spelling error, check GBIF"
-                    print(f"{species_name:<30} {'Not found':<20} {'Not found':<20} ERROR")
-                    print(f"Error: {error_msg}")
-                    sys.exit(1)  # Stop the script
-                    
-            except Exception as e:
-                error_msg = f"Error retrieving data for species '{species_name}': {str(e)}"
-                print(f"{species_name:<30} {'Error':<20} {'Error':<20} FAILED")
-                print(f"Error: {error_msg}")
-                sys.exit(1)  # Stop the script
-        
-        taxonomy[1] = sorted(list(set(taxonomy[1])))
-        print("-" * 80)
-        
-        num_orders = len(taxonomy[1])
-        num_families = len(taxonomy[2])
-        num_species = len(taxonomy[3])
-        
-        print("\nOrder indices:")
-        for i, order in enumerate(taxonomy[1]):
-            print(f"  {i}: {order}")
-        
-        print("\nFamily indices:")
-        for i, family in enumerate(taxonomy[2].keys()):
-            print(f"  {i}: {family}")
-        
-        print("\nSpecies indices:")
-        for i, species in enumerate(species_list):
-            print(f"  {i}: {species}")
-        
-        print(f"\nTaxonomy built: {num_orders} orders, {num_families} families, {num_species} species")
-        
-        return taxonomy, species_to_family, family_to_order
-    
-    def get_metrics(self, predicted_frames, true_frames, labels):
-        """
-        Calculate precision, recall, and F1 score for object detection results.
-        """
-        def calculate_iou(box1, box2):
-            x1_min, y1_min = box1[1] - box1[3] / 2, box1[2] - box1[4] / 2
-            x1_max, y1_max = box1[1] + box1[3] / 2, box1[2] + box1[4] / 2
-            x2_min, y2_min = box2[1] - box2[3] / 2, box2[2] - box2[4] / 2
-            x2_max, y2_max = box2[1] + box2[3] / 2, box2[2] + box2[4] / 2
-
-            inter_x_min = max(x1_min, x2_min)
-            inter_y_min = max(y1_min, y2_min)
-            inter_x_max = min(x1_max, x2_max)
-            inter_y_max = min(y1_max, y2_max)
-
-            inter_area = max(0, inter_x_max - inter_x_min) * max(0, inter_y_max - inter_y_min)
-            box1_area = (x1_max - x1_min) * (y1_max - y1_min)
-            box2_area = (x2_max - x2_min) * (y2_max - y2_min)
-
-            iou = inter_area / (box1_area + box2_area - inter_area)
-            return iou
-
-        def calculate_precision_recall(pred_boxes, true_boxes, iou_threshold=0.5):
-            label_results = defaultdict(lambda: {'tp': 0, 'fp': 0, 'fn': 0})
-            generic_tp = 0
-            generic_fp = 0
-            
-            matched_true_boxes = set()
-            
-            for pred_box in pred_boxes:
-                label_idx = pred_box[0]
-                matched = False
-                
-                best_iou = 0
-                best_match_idx = -1
-                
-                for i, true_box in enumerate(true_boxes):
-                    if i in matched_true_boxes:
-                        continue
-                    
-                    iou = calculate_iou(pred_box, true_box)
-                    if iou >= iou_threshold and iou > best_iou:
-                        best_iou = iou
-                        best_match_idx = i
-                
-                if best_match_idx >= 0:
-                    matched = True
-                    true_box = true_boxes[best_match_idx]
-                    matched_true_boxes.add(best_match_idx)
-                    generic_tp += 1
-                    
-                    if pred_box[0] == true_box[0]:
-                        label_results[label_idx]['tp'] += 1
-                    else:
-                        label_results[label_idx]['fp'] += 1
-                        true_label_idx = true_box[0]
-                        label_results[true_label_idx]['fn'] += 1
-                        
-                if not matched:
-                    label_results[label_idx]['fp'] += 1
-                    generic_fp += 1
-            
-            for i, true_box in enumerate(true_boxes):
-                if i not in matched_true_boxes:
-                    label_idx = true_box[0]
-                    label_results[label_idx]['fn'] += 1
-            
-            generic_fn = len(true_boxes) - len(matched_true_boxes)
-            
-            return label_results, generic_tp, generic_fp, generic_fn
-
-        label_metrics = defaultdict(lambda: {'tp': 0, 'fp': 0, 'fn': 0, 'support': 0})
-        background_metrics = {'tp': 0, 'fp': 0, 'fn': 0, 'support': 0}
-        generic_metrics = {'tp': 0, 'fp': 0, 'fn': 0}
-        
-        for true_frame in true_frames:
-            if not true_frame:  # Empty frame (background only)
-                background_metrics['support'] += 1
-            else:
-                for true_box in true_frame:
-                    label_idx = true_box[0]
-                    label_metrics[label_idx]['support'] += 1  # Count each detection, not just unique labels
-
-        for pred_frame, true_frame in zip(predicted_frames, true_frames):
-            if not pred_frame and not true_frame:
-                background_metrics['tp'] += 1
-            elif not pred_frame:
-                background_metrics['fn'] += 1
-            elif not true_frame:
-                background_metrics['fp'] += 1
-            else:
-                frame_results, g_tp, g_fp, g_fn = calculate_precision_recall(pred_frame, true_frame)
-                
-                for label_idx, metrics in frame_results.items():
-                    label_metrics[label_idx]['tp'] += metrics['tp']
-                    label_metrics[label_idx]['fp'] += metrics['fp'] 
-                    label_metrics[label_idx]['fn'] += metrics['fn']
-                
-                generic_metrics['tp'] += g_tp
-                generic_metrics['fp'] += g_fp
-                generic_metrics['fn'] += g_fn
-
-        table_data = []
-        # Store individual class metrics for macro-averaging
-        class_precisions = []
-        class_recalls = []
-        class_f1s = []
-
-        for label_idx, metrics in label_metrics.items():
-            tp = metrics['tp']
-            fp = metrics['fp']
-            fn = metrics['fn']
-            support = metrics['support']
-            precision = tp / (tp + fp) if (tp + fp) > 0 else 0
-            recall = tp / (tp + fn) if (tp + fn) > 0 else 0
-            f1_score = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0
-            
-            # Store for macro-averaging
-            class_precisions.append(precision)
-            class_recalls.append(recall)
-            class_f1s.append(f1_score)
-            
-            label_name = labels[label_idx] if label_idx < len(labels) else f"Label {label_idx}"
-            table_data.append([label_name, f"{precision:.2f}", f"{recall:.2f}", f"{f1_score:.2f}", f"{support}"])
-
-            print(f"Debug {label_name}: TP={tp}, FP={fp}, FN={fn}")
-            print(f"  Raw P={tp/(tp+fp) if (tp+fp)>0 else 0}, R={tp/(tp+fn) if (tp+fn)>0 else 0}")
-
-        tp = background_metrics['tp']
-        fp = background_metrics['fp']
-        fn = background_metrics['fn']
-        support = background_metrics['support']
-        precision = tp / (tp + fp) if (tp + fp) > 0 else 0
-        recall = tp / (tp + fn) if (tp + fn) > 0 else 0
-        f1_score = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0
-        table_data.append(["Background", f"{precision:.2f}", f"{recall:.2f}", f"{f1_score:.2f}", f"{support}"])
-
-        headers = ["Label", "Precision", "Recall", "F1 Score", "Support"]
-        total_tp = sum(metrics['tp'] for metrics in label_metrics.values())
-        total_fp = sum(metrics['fp'] for metrics in label_metrics.values())
-        total_fn = sum(metrics['fn'] for metrics in label_metrics.values())
-        total_support = sum(metrics['support'] for metrics in label_metrics.values())
-
-        total_precision = total_tp / (total_tp + total_fp) if (total_tp + total_fp) > 0 else 0
-        total_recall = total_tp / (total_tp + total_fn) if (total_tp + total_fn) > 0 else 0
-        total_f1_score = 2 * (total_precision * total_recall) / (total_precision + total_recall) if (total_precision + total_recall) > 0 else 0
-
-        table_data.append(["\nTotal (micro-avg, excl. background)", f"{total_precision:.2f}", f"{total_recall:.2f}", f"{total_f1_score:.2f}", f"{total_support}"])
-        
-        # Add macro-average
-        if class_precisions:
-            macro_precision = sum(class_precisions) / len(class_precisions)
-            macro_recall = sum(class_recalls) / len(class_recalls)
-            macro_f1 = sum(class_f1s) / len(class_f1s)
-            table_data.append(["Total (macro-avg, excl. background)", f"{macro_precision:.2f}", f"{macro_recall:.2f}", f"{macro_f1:.2f}", f"{total_support}"])
-        
-        print(tabulate(table_data, headers=headers, tablefmt="grid"))
-        
-        generic_tp = generic_metrics['tp']
-        generic_fp = generic_metrics['fp']
-        generic_fn = generic_metrics['fn']
-
-        generic_precision = generic_tp / (generic_tp + generic_fp) if (generic_tp + generic_fp) > 0 else 0
-        generic_recall = generic_tp / (generic_tp + generic_fn) if (generic_tp + generic_fn) > 0 else 0
-        generic_f1_score = 2 * (generic_precision * generic_recall) / (generic_precision + generic_recall) if (generic_precision + generic_recall) > 0 else 0
-
-        print("\nGeneric Total", f"{generic_precision:.2f}", f"{generic_recall:.2f}", f"{generic_f1_score:.2f}")
-        
-        return total_precision, total_recall, total_f1_score
-
-    def run(self, test_dir):
-        predicted_frames, true_frames, total_time = self.get_frames(test_dir)
-        num_frames = len(os.listdir(os.path.join(test_dir, 'images')))
-        avg_time_per_frame = total_time / num_frames
-
-        print(f"\nTotal time: {total_time:.2f} seconds")
-        print(f"Average time per frame: {avg_time_per_frame:.4f} seconds")
-        
-        # Get taxonomy information for hierarchical analysis
-        taxonomy, species_to_family, family_to_order = self.get_taxonomic_info(self.species_names)
-        family_list = list(family_to_order.keys())
-        order_list = list(taxonomy[1])
-        
-        # Convert species-level predictions to family and order levels
-        true_family_frames = []
-        true_order_frames = []
-        predicted_family_frames = []
-        predicted_order_frames = []
-        
-        for true_frame in true_frames:
-            frame_family_boxes = []
-            frame_order_boxes = []
-            
-            if true_frame:
-                for true_box in true_frame:
-                    species_idx = true_box[0]
-                    species_name = self.species_names[species_idx]
-                    family_name = species_to_family[species_name]
-                    order_name = family_to_order[family_name]
-                    
-                    family_label = [family_list.index(family_name)] + list(true_box[1:])
-                    order_label = [order_list.index(order_name)] + list(true_box[1:])
-                    
-                    frame_family_boxes.append(family_label)
-                    frame_order_boxes.append(order_label)
-                    
-            true_family_frames.append(frame_family_boxes)
-            true_order_frames.append(frame_order_boxes)
-        
-        for pred_frame in predicted_frames:
-            frame_family_boxes = []
-            frame_order_boxes = []
-            
-            if pred_frame:
-                for pred_box in pred_frame:
-                    species_idx = pred_box[0]
-                    species_name = self.species_names[species_idx]
-                    family_name = species_to_family[species_name]
-                    order_name = family_to_order[family_name]
-                    
-                    family_label = [family_list.index(family_name)] + list(map(np.float32, pred_box[1:]))
-                    order_label = [order_list.index(order_name)] + list(map(np.float32, pred_box[1:]))
-                    
-                    frame_family_boxes.append(family_label)
-                    frame_order_boxes.append(order_label)
-                    
-            predicted_family_frames.append(frame_family_boxes)
-            predicted_order_frames.append(frame_order_boxes)
-        
-        # Display metrics for all taxonomic levels
-        print("\nSpecies Level Metrics")
-        self.get_metrics(predicted_frames, true_frames, self.species_names)
-        
-        print("\nFamily Level Metrics")
-        self.get_metrics(predicted_family_frames, true_family_frames, family_list)
-        
-        print("\nOrder Level Metrics")
-        self.get_metrics(predicted_order_frames, true_order_frames, order_list)
-
-if __name__ == "__main__":
-    species_names = [
-        "Coccinella septempunctata", "Apis mellifera", "Bombus lapidarius", "Bombus terrestris",
-        "Eupeodes corollae", "Episyrphus balteatus", "Aglais urticae", "Vespula vulgaris",
-        "Eristalis tenax"
-    ]
-    
-    test_resnet(
-        data_path="/mnt/nvme0n1p1/mit/two-stage-detection/bjerge-test",
-        yolo_weights="/mnt/nvme0n1p1/mit/two-stage-detection/small-generic.pt",
-        resnet_weights="/mnt/nvme0n1p1/mit/two-stage-detection/output/best_resnet50.pt",
-        model="resnet50",
-        species_names=species_names,
-        output_dir="/mnt/nvme0n1p1/mit/two-stage-detection/output"
-    )