add lab 2

.gitattributes

@@ -0,0 +1 @@
+*.pth filter=lfs diff=lfs merge=lfs -text

lab02_Gun_detection_fasterRCNN/gun_detection_fasterRCNN.ipynb

@@ -0,0 +1,176 @@
+# Gun Detection Lab: Detecting Guns in Images using PyTorch and Faster R-CNN
+
+# Step 1: Install Dependencies
+!pip install torch torchvision matplotlib opencv-python pycocotools py7zr
+
+# Step 2: Import Libraries
+import os
+import torch
+import torchvision
+from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
+from torchvision.transforms import functional as F
+from torch.utils.data import DataLoader
+from PIL import Image
+import matplotlib.pyplot as plt
+import numpy as np
+import py7zr
+
+# Step 3: Download and Extract Dataset
+dataset_url = "https://github.com/frankwxu/AI4DigitalForensics/raw/main/lab02_Gun_detection_fasterRCNN/data/data.7z"
+dataset_path = "/content/data.7z"
+extracted_path = "/content/dataset"
+
+# Download the dataset
+!wget -O {dataset_path} {dataset_url}
+
+# Extract the dataset
+with py7zr.SevenZipFile(dataset_path, mode='r') as z:
+    z.extractall(path=extracted_path)
+
+# Define paths to images and labels
+images_dir = os.path.join(extracted_path, "Images")  # Folder name is "Images"
+labels_dir = os.path.join(extracted_path, "Labels")  # Folder name is "Labels"
+
+# Step 4: Create Custom Dataset Class
+class GunDataset(torch.utils.data.Dataset):
+    def __init__(self, root, transforms=None):
+        self.root = root
+        self.transforms = transforms
+        self.imgs = list(sorted(os.listdir(os.path.join(root, "Images"))))
+        self.labels = list(sorted(os.listdir(os.path.join(root, "Labels"))))
+    
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.root, "Images", self.imgs[idx])
+        label_path = os.path.join(self.root, "Labels", self.labels[idx])
+        
+        img = Image.open(img_path).convert("RGB")
+        target = self.parse_annotation(label_path)
+        
+        if self.transforms is not None:
+            img = self.transforms(img)
+        
+        return img, target
+    
+    def __len__(self):
+        return len(self.imgs)
+    
+    def parse_annotation(self, label_path):
+        with open(label_path, 'r') as f:
+            lines = f.readlines()
+            num_guns = int(lines[0].strip())  # Number of guns in the image
+            bboxes = []
+            labels = []
+            
+            # Parse each bounding box
+            for i in range(1, num_guns + 1):
+                bbox = list(map(int, lines[i].strip().split()))  # Bounding box (xmin, ymin, xmax, ymax)
+                bboxes.append(bbox)
+                labels.append(1)  # All objects are labeled as "Gun" (class 1)
+        
+        return {
+            'boxes': torch.tensor(bboxes, dtype=torch.float32),
+            'labels': torch.tensor(labels, dtype=torch.int64)
+        }
+
+# Step 5: Define Data Transforms
+def get_transform(train):
+    transforms = []
+    transforms.append(torchvision.transforms.ToTensor())
+    if train:
+        transforms.append(torchvision.transforms.RandomHorizontalFlip(0.5))
+    return torchvision.transforms.Compose(transforms)
+
+# Step 6: Load Dataset
+train_dataset = GunDataset(extracted_path, get_transform(train=True))
+test_dataset = GunDataset(extracted_path, get_transform(train=False))
+
+train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True, collate_fn=lambda x: tuple(zip(*x)))
+test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False, collate_fn=lambda x: tuple(zip(*x)))
+
+# Step 7: Load Pre-trained Faster R-CNN Model
+model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
+
+# Replace the classifier head for custom classes
+num_classes = 2  # Background + Gun
+in_features = model.roi_heads.box_predictor.cls_score.in_features
+model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
+
+# Move the model to GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+# Step 8: Define Optimizer
+params = [p for p in model.parameters() if p.requires_grad]
+optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)
+
+# Step 9: Train the Model
+def train_one_epoch(model, optimizer, data_loader, device, epoch):
+    model.train()
+    for i, (images, targets) in enumerate(data_loader):
+        images = list(image.to(device) for image in images)
+        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
+        
+        loss_dict = model(images, targets)
+        losses = sum(loss for loss in loss_dict.values())
+        
+        optimizer.zero_grad()
+        losses.backward()
+        optimizer.step()
+        
+        if i % 10 == 0:
+            print(f"Iteration {i}, Loss: {losses.item()}")
+
+# Train for 10 epochs
+num_epochs = 10
+for epoch in range(num_epochs):
+    print(f"Epoch {epoch+1}/{num_epochs}")
+    train_one_epoch(model, optimizer, train_loader, device, epoch)
+
+# Step 10: Evaluate the Model
+def evaluate_model(model, data_loader, device):
+    model.eval()
+    all_preds = []
+    all_labels = []
+    
+    with torch.no_grad():
+        for images, targets in data_loader:
+            images = list(image.to(device) for image in images)
+            outputs = model(images)
+            
+            for output in outputs:
+                preds = output['labels'].cpu().numpy()
+                all_preds.extend(preds)
+            
+            for target in targets:
+                labels = target['labels'].cpu().numpy()
+                all_labels.extend(labels)
+    
+    print("\nClassification Report:")
+    from sklearn.metrics import classification_report
+    print(classification_report(all_labels, all_preds, target_names=['Background', 'Gun']))
+
+evaluate_model(model, test_loader, device)
+
+# Step 11: Test on Custom Images
+def predict_image(img_path, model, transform):
+    img = Image.open(img_path).convert("RGB")
+    img_tensor = transform(img).unsqueeze(0).to(device)
+    
+    model.eval()
+    with torch.no_grad():
+        prediction = model(img_tensor)[0]
+    
+    # Plot the image with bounding boxes
+    plt.figure(figsize=(10, 10))
+    plt.imshow(img)
+    for box, label in zip(prediction['boxes'], prediction['labels']):
+        if label == 1:  # Gun detected
+            box = box.cpu().numpy()
+            plt.gca().add_patch(plt.Rectangle((box[0], box[1]), box[2]-box[0], box[3]-box[1], 
+                                              edgecolor='red', facecolor='none', linewidth=2))
+    plt.axis('off')
+    plt.show()
+
+# Example Usage
+custom_image_path = '/content/dataset/Images/test_image.jpg'  # Update this path
+predict_image(custom_image_path, model, get_transform(train=False))