diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..ec4a626 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +*.pth filter=lfs diff=lfs merge=lfs -text diff --git a/lab02_Gun_detection_fasterRCNN/gun_detection_fasterRCNN.ipynb b/lab02_Gun_detection_fasterRCNN/gun_detection_fasterRCNN.ipynb new file mode 100644 index 0000000..88ef2d9 --- /dev/null +++ b/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)) \ No newline at end of file diff --git a/lab02_Gun_detection_fasterRCNN/labxx.ipynb b/lab02_Gun_detection_fasterRCNN/labxx.ipynb deleted file mode 100644 index e69de29..0000000