changed api structure and added object detection
This commit is contained in:
@@ -0,0 +1,152 @@
|
||||
import os
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from tqdm import tqdm
|
||||
|
||||
class SeparableConvolution(nn.Module):
|
||||
def __init__(self, c_in, c_out, kernel_size):
|
||||
super().__init__()
|
||||
self.depthwise = nn.Conv2d(c_in, c_in, kernel_size, groups=c_in, padding=kernel_size//2)
|
||||
self.bn1 = nn.BatchNorm2d(c_in)
|
||||
self.pointwise = nn.Conv2d(c_in, c_out, kernel_size=1)
|
||||
self.bn2 = nn.BatchNorm2d(c_out)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.depthwise(x)
|
||||
x = self.bn1(x)
|
||||
x = F.relu(x)
|
||||
|
||||
x = self.pointwise(x)
|
||||
x = self.bn2(x)
|
||||
x = F.relu(x)
|
||||
|
||||
return x
|
||||
|
||||
class SkipBlock(nn.Module):
|
||||
def __init__(self, c_in, c_out, kernel_size=3):
|
||||
super().__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(c_in, c_out, kernel_size, padding=kernel_size//2),
|
||||
nn.BatchNorm2d(c_out),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(c_out, c_out, kernel_size, padding=kernel_size//2),
|
||||
nn.BatchNorm2d(c_out),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(c_out, c_out, kernel_size, padding=kernel_size//2),
|
||||
nn.BatchNorm2d(c_out),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
self.conv_skip = nn.Sequential(
|
||||
nn.Conv2d(c_in, c_out, 1),
|
||||
nn.BatchNorm2d(c_out),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return(F.relu(self.conv_skip(x) + self.conv(x), inplace=True))
|
||||
|
||||
|
||||
class Bird_CNN(nn.Module):
|
||||
def __init__(self, c_in, c_hidden, c_out):
|
||||
super().__init__()
|
||||
|
||||
self.model = nn.Sequential(
|
||||
nn.Conv2d(c_in, c_hidden, kernel_size=3, padding=1),
|
||||
nn.BatchNorm2d(c_hidden),
|
||||
nn.ReLU(inplace=True),
|
||||
|
||||
SkipBlock(c_in=c_hidden, c_out=c_hidden),
|
||||
SkipBlock(c_in=c_hidden, c_out=c_hidden),
|
||||
SkipBlock(c_in=c_hidden, c_out=c_hidden),
|
||||
SkipBlock(c_in=c_hidden, c_out=c_hidden),
|
||||
|
||||
SkipBlock(c_in=c_hidden, c_out=c_hidden*2),
|
||||
SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
|
||||
SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
|
||||
SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
|
||||
|
||||
nn.Conv2d(c_hidden*2, c_hidden*4, kernel_size=3, padding=1),
|
||||
nn.ReLU(inplace=True),
|
||||
|
||||
nn.AdaptiveAvgPool2d((1, 1)),
|
||||
nn.Flatten(),
|
||||
nn.Linear(c_hidden*4, c_out),
|
||||
nn.Dropout(0.3)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.model(x)
|
||||
|
||||
|
||||
def trainCNN(model, optimizer, loss_module, train_data_loader, validation_data_loader, device, num_epochs, SAVE_PATH, save=False):
|
||||
|
||||
best_val = 0
|
||||
for epoch in range(num_epochs):
|
||||
############
|
||||
# Training #
|
||||
############
|
||||
model.train()
|
||||
|
||||
true_preds, count = 0, 0
|
||||
for data_inputs, classes in tqdm(train_data_loader, desc=f"Train Epoch {epoch+1}", leave=False):
|
||||
data_inputs = data_inputs.to(device)
|
||||
classes = classes.to(device)
|
||||
|
||||
preds = model(data_inputs)
|
||||
|
||||
loss = loss_module(preds, classes)
|
||||
|
||||
optimizer.zero_grad()
|
||||
|
||||
loss.backward()
|
||||
|
||||
optimizer.step()
|
||||
|
||||
true_preds += (preds.argmax(dim=1) == classes).sum().item()
|
||||
count += data_inputs.size(0)
|
||||
train_acc = true_preds / count
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
##############
|
||||
# Validation #
|
||||
##############
|
||||
model.eval()
|
||||
|
||||
true_preds, count = 0, 0
|
||||
for data_inputs, classes in tqdm(validation_data_loader, desc=f"Validate Epoch {epoch+1}", leave=False):
|
||||
with torch.no_grad():
|
||||
data_inputs = data_inputs.to(device)
|
||||
classes = classes.to(device)
|
||||
|
||||
preds = model(data_inputs)
|
||||
|
||||
loss = loss_module(preds, classes)
|
||||
|
||||
true_preds += (preds.argmax(dim=1) == classes).sum().item()
|
||||
count += data_inputs.size(0)
|
||||
val_acc = true_preds / count
|
||||
|
||||
if(save and best_val < val_acc):
|
||||
best_val = val_acc
|
||||
save_dir = os.path.join(SAVE_PATH, "bird_cnn")
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
save_path = os.path.join(save_dir, f"bird_cnn")
|
||||
torch.save(model.state_dict(), save_path)
|
||||
|
||||
print(f"epoch: {epoch+1} | train accuracy: {int(train_acc * 1000) / 10}% | validation accuracy: {int(val_acc * 1000) / 10}%")
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def sample(model, img, device, SAVE_PATH, model_name="bird_cnn", folder="bird_cnn"):
|
||||
with torch.no_grad():
|
||||
full_path = os.path.join(SAVE_PATH, folder, model_name)
|
||||
state_dict = torch.load(full_path, weights_only=False)
|
||||
model.load_state_dict(state_dict)
|
||||
model.eval()
|
||||
img = img.to(device)
|
||||
pred = model(img)
|
||||
probs = F.softmax(pred, dim=1)
|
||||
return(torch.max(probs, dim=1))
|
||||
@@ -0,0 +1,74 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision
|
||||
from torchvision import datasets, transforms
|
||||
from torch.utils.data import DataLoader, random_split
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
from util import TransformedSubset, visualizeData
|
||||
from .bird_cnn import Bird_CNN, sample, trainCNN
|
||||
|
||||
from enum import Enum
|
||||
|
||||
from PIL import Image
|
||||
|
||||
SAVE_PATH = "../saved_models"
|
||||
IMAGE_SIZE = 128
|
||||
|
||||
class bird_species(Enum):
|
||||
Common_Kingfisher = 0
|
||||
Common_Myna = 1
|
||||
House_Crow = 2
|
||||
Indian_Peacock = 3
|
||||
Indian_Pitta = 4
|
||||
Ruddy_Shelduck = 5
|
||||
Sarus_Crane = 6
|
||||
|
||||
transform_augemnt = transforms.Compose([
|
||||
transforms.RandomAffine(
|
||||
degrees=35, # no rotation
|
||||
translate=(0.2, 0.2) # shift up to 20% horizontally/vertically
|
||||
),
|
||||
transforms.RandomHorizontalFlip(p=0.5),
|
||||
transforms.Resize(IMAGE_SIZE),
|
||||
transforms.CenterCrop(IMAGE_SIZE),
|
||||
transforms.ToTensor()
|
||||
])
|
||||
|
||||
transform = transforms.Compose([
|
||||
transforms.Resize(IMAGE_SIZE),
|
||||
transforms.CenterCrop(IMAGE_SIZE),
|
||||
transforms.ToTensor()
|
||||
])
|
||||
|
||||
dataset = datasets.ImageFolder("data/CUB_200_2011/images")
|
||||
|
||||
train_size = int(0.8 * len(dataset))
|
||||
val_size = len(dataset) - train_size
|
||||
|
||||
train_subset, val_subset = random_split(dataset, [train_size, val_size])
|
||||
train_dataset = TransformedSubset(train_subset, transform_augemnt)
|
||||
val_dataset = TransformedSubset(val_subset, transform)
|
||||
|
||||
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
|
||||
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
|
||||
|
||||
device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
|
||||
print("Using device", device)
|
||||
|
||||
model = Bird_CNN(c_in=3, c_hidden=16, c_out=200)
|
||||
model.to(device)
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
|
||||
loss_module = nn.CrossEntropyLoss()
|
||||
|
||||
trainCNN(model, optimizer, loss_module, train_loader, val_loader, device, 500, SAVE_PATH=SAVE_PATH, save=True)
|
||||
exit()
|
||||
|
||||
image_path = "test1.jpg"
|
||||
|
||||
image = Image.open(image_path).convert("RGB")
|
||||
image = transform(image)
|
||||
image = image.unsqueeze(0)
|
||||
|
||||
confidence, pred = sample(model=model, img=image, device=device, SAVE_PATH=SAVE_PATH)
|
||||
print(f"Species: {bird_species(pred.item()).name} | Confidence: {int(confidence.item()*100)/100}")
|
||||
@@ -0,0 +1,125 @@
|
||||
import os
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torch.utils.data import Dataset
|
||||
from pycocotools.coco import COCO
|
||||
|
||||
class CocoDetectionDataset(Dataset):
|
||||
# Init function: loads annotation file and prepares list of image IDs
|
||||
def __init__(self, image_dir, annotation_path, transforms=None):
|
||||
self.image_dir = image_dir
|
||||
self.coco = COCO(annotation_path)
|
||||
self.image_ids = list(self.coco.imgs.keys())
|
||||
self.transforms = transforms
|
||||
|
||||
# Returns total number of images
|
||||
def __len__(self):
|
||||
return len(self.image_ids)
|
||||
|
||||
# Fetches a single image and its annotations
|
||||
def __getitem__(self, idx):
|
||||
image_id = self.image_ids[idx]
|
||||
image_info = self.coco.loadImgs(image_id)[0]
|
||||
image_path = os.path.join(self.image_dir, image_info['file_name'])
|
||||
image = Image.open(image_path).convert("RGB")
|
||||
|
||||
# Load all annotations for this image
|
||||
annotation_ids = self.coco.getAnnIds(imgIds=image_id)
|
||||
annotations = self.coco.loadAnns(annotation_ids)
|
||||
|
||||
# Extract bounding boxes and labels from annotations
|
||||
boxes = []
|
||||
labels = []
|
||||
for obj in annotations:
|
||||
xmin, ymin, width, height = obj['bbox']
|
||||
xmin, ymin, width, height = float(xmin), float(ymin), float(width), float(height)
|
||||
xmax = xmin + width
|
||||
ymax = ymin + height
|
||||
boxes.append([xmin, ymin, xmax, ymax])
|
||||
labels.append(obj['category_id'])
|
||||
|
||||
# Convert annotations to PyTorch tensors
|
||||
boxes = torch.as_tensor(boxes, dtype=torch.float32)
|
||||
labels = torch.as_tensor(labels, dtype=torch.int64)
|
||||
area = torch.as_tensor([obj['area'] for obj in annotations], dtype=torch.float32)
|
||||
iscrowd = torch.as_tensor([obj.get('iscrowd', 0) for obj in annotations], dtype=torch.int64)
|
||||
|
||||
# Package everything into a target dictionary
|
||||
target = {
|
||||
"boxes": boxes,
|
||||
"labels": labels,
|
||||
"image_id": image_id,
|
||||
"area": area,
|
||||
"iscrowd": iscrowd
|
||||
}
|
||||
|
||||
# Apply transforms if any were passed
|
||||
if self.transforms:
|
||||
image = self.transforms(image)
|
||||
|
||||
return image, target
|
||||
|
||||
|
||||
class CocoDetectionDatasetResized(Dataset):
|
||||
def __init__(self, image_dir, annotation_path, img_size=64, transforms=None):
|
||||
self.image_dir = image_dir
|
||||
self.coco = COCO(annotation_path)
|
||||
self.image_ids = list(self.coco.imgs.keys())
|
||||
self.transforms = transforms
|
||||
self.img_size = img_size
|
||||
|
||||
def __len__(self):
|
||||
return len(self.image_ids)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
image_id = self.image_ids[idx]
|
||||
image_info = self.coco.loadImgs(image_id)[0]
|
||||
image_path = os.path.join(self.image_dir, image_info['file_name'])
|
||||
|
||||
image = Image.open(image_path).convert("RGB")
|
||||
|
||||
orig_w, orig_h = image.size
|
||||
|
||||
scale_w = self.img_size / orig_w
|
||||
scale_h = self.img_size / orig_h
|
||||
|
||||
# Load annotations
|
||||
annotation_ids = self.coco.getAnnIds(imgIds=image_id)
|
||||
annotations = self.coco.loadAnns(annotation_ids)
|
||||
|
||||
boxes = []
|
||||
labels = []
|
||||
|
||||
for obj in annotations:
|
||||
xmin, ymin, width, height = obj['bbox']
|
||||
xmin, ymin, width, height = float(xmin), float(ymin), float(width), float(height)
|
||||
|
||||
xmin = xmin * scale_w
|
||||
ymin = ymin * scale_h
|
||||
xmax = (xmin + width * scale_w)
|
||||
ymax = (ymin + height * scale_h)
|
||||
|
||||
boxes.append([xmin, ymin, xmax, ymax])
|
||||
labels.append(obj['category_id'])
|
||||
|
||||
boxes = torch.tensor(boxes, dtype=torch.float32)
|
||||
labels = torch.tensor(labels, dtype=torch.long)
|
||||
|
||||
area = torch.tensor([obj['area'] * scale_w * scale_h for obj in annotations], dtype=torch.float32)
|
||||
iscrowd = torch.tensor([obj.get('iscrowd', 0) for obj in annotations], dtype=torch.long)
|
||||
|
||||
target = {
|
||||
"boxes": boxes,
|
||||
"labels": labels,
|
||||
"image_id": torch.tensor([image_id]),
|
||||
"area": area,
|
||||
"iscrowd": iscrowd
|
||||
}
|
||||
|
||||
# resize image
|
||||
image = image.resize((self.img_size, self.img_size))
|
||||
|
||||
if self.transforms:
|
||||
image = self.transforms(image)
|
||||
|
||||
return image, target
|
||||
@@ -0,0 +1,118 @@
|
||||
from pathlib import Path
|
||||
from tqdm import tqdm
|
||||
from ..util import iou
|
||||
from .cocoDetectionDataset import CocoDetectionDataset
|
||||
from torchvision.transforms import ToPILImage, ToTensor
|
||||
import torch
|
||||
import torchvision.transforms.functional as TF
|
||||
from torch.utils.data import DataLoader
|
||||
import random
|
||||
|
||||
|
||||
|
||||
SAVE_PATH = "./saved_models"
|
||||
|
||||
def get_transform():
|
||||
return ToTensor()
|
||||
|
||||
dataset = CocoDetectionDataset(
|
||||
image_dir="data/faces/train",
|
||||
annotation_path="data/faces/train/_annotations.coco.json",
|
||||
transforms=get_transform()
|
||||
)
|
||||
|
||||
processing_loader = DataLoader(dataset, batch_size=1, shuffle=False, collate_fn=lambda x: tuple(zip(*x)))
|
||||
|
||||
def random_crop(W, H, sizeX=64, sizeY =64):
|
||||
x = random.randint(0, max(W - sizeX, 0))
|
||||
y = random.randint(0, max(H - sizeY, 0))
|
||||
return x, y, x + sizeX, y + sizeY
|
||||
|
||||
def box_inside(inner_box, outer_box):
|
||||
return (
|
||||
inner_box[0] >= outer_box[0] and
|
||||
inner_box[1] >= outer_box[1] and
|
||||
inner_box[2] <= outer_box[2] and
|
||||
inner_box[3] <= outer_box[3]
|
||||
)
|
||||
|
||||
def is_background(crop, gt_boxes, threshold=0.0):
|
||||
for box in gt_boxes:
|
||||
if iou(crop, box) > threshold:
|
||||
return False
|
||||
return True
|
||||
|
||||
def get_background_crop(image, gt_boxes, max_trials=100):
|
||||
C, H, W = image.shape
|
||||
|
||||
for _ in range(max_trials):
|
||||
x1, y1, x2, y2 = random_crop(W, H, sizeX=random.randint(50, 600), sizeY=random.randint(50, 600))
|
||||
|
||||
crop_box = (x1, y1, x2, y2)
|
||||
|
||||
if is_background(crop_box, gt_boxes, threshold=0.0):
|
||||
crop = image[:, y1:y2, x1:x2]
|
||||
return crop
|
||||
|
||||
return None
|
||||
|
||||
def create_background_tensor(amount, dataset, labels, boxes, image):
|
||||
for _ in range(amount):
|
||||
background = get_background_crop(image=image, gt_boxes=boxes.to(torch.int64))
|
||||
if background is None:
|
||||
break
|
||||
background = TF.resize(background, [64, 64], antialias=True)
|
||||
dataset.append(background)
|
||||
labels.append(torch.tensor(0))
|
||||
return dataset, labels
|
||||
|
||||
def create_stack(images, annotations, PADDING):
|
||||
croped_images = []
|
||||
labels = []
|
||||
for i in range(len(images)):
|
||||
image = images[i]
|
||||
annotation = annotations[i]
|
||||
_, h, w = image.shape
|
||||
|
||||
for box, label in tuple(zip(annotation["boxes"], annotation["labels"])):
|
||||
box = box.to(torch.int64)
|
||||
box_copy = []
|
||||
|
||||
box_copy.append(max(box[0]-PADDING, 0))
|
||||
box_copy.append(max(box[1]-PADDING, 0))
|
||||
box_copy.append(min(box[2]+PADDING, w-1))
|
||||
box_copy.append(min(box[3]+PADDING, h-1))
|
||||
|
||||
if box_copy[1] == box_copy[3] or box_copy[0] == box_copy[2]:
|
||||
continue
|
||||
|
||||
croped_image = image[: , box_copy[1]:box_copy[3], box_copy[0]:box_copy[2]]
|
||||
croped_image = TF.resize(croped_image, [64, 64], antialias=True)
|
||||
croped_images.append(croped_image)
|
||||
labels.append(label)
|
||||
|
||||
croped_images, labels = create_background_tensor(amount=10, dataset=croped_images, labels=labels, boxes=annotation["boxes"], image=image)
|
||||
return croped_images, labels
|
||||
|
||||
def create_croped_dataset():
|
||||
to_pil = ToPILImage()
|
||||
|
||||
out_path = Path("data/faces_processed/face")
|
||||
out_path.mkdir(parents=True, exist_ok=True)
|
||||
out_path = Path("data/faces_processed/background")
|
||||
out_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
count_image = 0
|
||||
count_background = 0
|
||||
for images, annotations in tqdm(processing_loader):
|
||||
croped_images, labels = create_stack(images, annotations, PADDING=0)
|
||||
|
||||
for i in range(len(croped_images)):
|
||||
pil_img = to_pil(croped_images[i])
|
||||
|
||||
if labels[i] == 1:
|
||||
pil_img.save(f"data/faces_processed/face/{count_image}.jpg")
|
||||
count_image += 1
|
||||
else:
|
||||
pil_img.save(f"data/faces_processed/background/{count_background}.jpg")
|
||||
count_background += 1
|
||||
@@ -0,0 +1,313 @@
|
||||
import os
|
||||
import cv2
|
||||
import torch
|
||||
import torchvision.transforms.functional as TF
|
||||
from tqdm import tqdm
|
||||
import torch.nn as nn
|
||||
from src.util import iou, visualizeImage
|
||||
import random
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class ObjectDetectionCNN(nn.Module):
|
||||
def __init__(self, c_in, c_hidden, c_out, layers):
|
||||
super().__init__()
|
||||
|
||||
self.model = nn.ModuleList()
|
||||
|
||||
self.model.append(nn.Sequential(
|
||||
nn.Conv2d(c_in, c_hidden, kernel_size=3, padding=1),
|
||||
nn.BatchNorm2d(c_hidden),
|
||||
nn.ReLU(inplace=True)
|
||||
))
|
||||
|
||||
for _ in range(layers-1):
|
||||
self.model.append(nn.Sequential(
|
||||
nn.Conv2d(c_hidden, c_hidden, kernel_size=3, padding=1),
|
||||
nn.BatchNorm2d(c_hidden),
|
||||
nn.ReLU(inplace=True)
|
||||
))
|
||||
|
||||
self.model.append(nn.Sequential(
|
||||
nn.AdaptiveAvgPool2d((1, 1)),
|
||||
nn.Flatten(),
|
||||
nn.Linear(c_hidden, c_out),
|
||||
nn.Dropout(0.3)
|
||||
))
|
||||
|
||||
def forward(self, x):
|
||||
for layer in self.model:
|
||||
x = layer(x)
|
||||
return x
|
||||
|
||||
|
||||
|
||||
def random_crop(W, H, sizeX=64, sizeY =64):
|
||||
x = random.randint(0, max(W - sizeX, 0))
|
||||
y = random.randint(0, max(H - sizeY, 0))
|
||||
return x, y, x + sizeX, y + sizeY
|
||||
|
||||
def is_background(crop, gt_boxes, threshold=0.0):
|
||||
for box in gt_boxes:
|
||||
if iou(crop, box) > threshold:
|
||||
return False
|
||||
return True
|
||||
|
||||
def get_background_crop(image, gt_boxes, max_trials=100):
|
||||
C, H, W = image.shape
|
||||
|
||||
for _ in range(max_trials):
|
||||
x1, y1, x2, y2 = random_crop(W, H, sizeX=random.randint(50, 600), sizeY=random.randint(50, 600))
|
||||
|
||||
crop_box = (x1, y1, x2, y2)
|
||||
|
||||
if is_background(crop_box, gt_boxes, threshold=0.1):
|
||||
crop = image[:, y1:y2, x1:x2]
|
||||
return crop
|
||||
|
||||
return None
|
||||
|
||||
def create_background_tensor(amount, dataset, labels, boxes, image):
|
||||
for _ in range(amount):
|
||||
background = get_background_crop(image=image, gt_boxes=boxes.to(torch.int64))
|
||||
if background is not None:
|
||||
background = TF.resize(background, [64, 64], antialias=True)
|
||||
dataset.append(background)
|
||||
labels.append(torch.tensor(0))
|
||||
else:
|
||||
print("Background not found")
|
||||
return dataset, labels
|
||||
|
||||
def create_stack(images, annotations, PADDING):
|
||||
croped_images = []
|
||||
labels = []
|
||||
for i in range(len(images)):
|
||||
image = images[i]
|
||||
annotation = annotations[i]
|
||||
_, h, w = image.shape
|
||||
|
||||
for box, label in tuple(zip(annotation["boxes"], annotation["labels"])):
|
||||
box = box.to(torch.int64)
|
||||
box_copy = []
|
||||
|
||||
box_copy.append(max(box[0]-PADDING, 0))
|
||||
box_copy.append(max(box[1]-PADDING, 0))
|
||||
box_copy.append(min(box[2]+PADDING, w-1))
|
||||
box_copy.append(min(box[3]+PADDING, h-1))
|
||||
|
||||
|
||||
croped_image = image[: , box_copy[1]:box_copy[3], box_copy[0]:box_copy[2]]
|
||||
#visualizeImage(croped_image)
|
||||
croped_image = TF.resize(croped_image, [64, 64], antialias=True)
|
||||
#visualizeImage(croped_image)
|
||||
croped_images.append(croped_image)
|
||||
labels.append(label)
|
||||
|
||||
#croped_images, labels = create_background_tensor(amount=len(croped_images), dataset=croped_images, labels=labels, boxes=annotation["boxes"], image=image)
|
||||
return croped_images, labels
|
||||
|
||||
def train(model, loss_module, train_loader, val_loader, optimizer, SAVE_PATH, model_name, saving=True, PADDING=20, device="cpu"):
|
||||
best_val = torch.finfo(torch.float32).max
|
||||
|
||||
for epoch in range(200):
|
||||
############
|
||||
# Training #
|
||||
############
|
||||
model.train()
|
||||
|
||||
true_preds, count, lossCount = 0, 0, 0.
|
||||
for images, annotations in tqdm(train_loader, desc=f"Train", leave=False):
|
||||
croped_images, labels = create_stack(images, annotations, PADDING)
|
||||
|
||||
croped_images = torch.stack(croped_images).to(device)
|
||||
labels = torch.stack(labels).to(device)
|
||||
|
||||
prediction = model(croped_images)
|
||||
|
||||
loss = loss_module(prediction, labels)
|
||||
lossCount += loss.sum().item()
|
||||
|
||||
optimizer.zero_grad()
|
||||
|
||||
loss.backward()
|
||||
|
||||
optimizer.step()
|
||||
|
||||
true_preds += (prediction.argmax(dim=1) == labels).sum().item()
|
||||
count += croped_images.size(0)
|
||||
|
||||
train_acc = true_preds / count
|
||||
train_loss = lossCount / count
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
##############
|
||||
# Validation #
|
||||
##############
|
||||
model.eval()
|
||||
|
||||
true_preds, count, lossCount = 0, 0, 0.
|
||||
for images, annotations in tqdm(val_loader, desc=f"Test", leave=False):
|
||||
with torch.no_grad():
|
||||
croped_images, labels = create_stack(images, annotations, PADDING)
|
||||
|
||||
croped_images = torch.stack(croped_images).to(device)
|
||||
labels = torch.stack(labels).to(device)
|
||||
|
||||
prediction = model(croped_images)
|
||||
|
||||
loss = loss_module(prediction, labels)
|
||||
lossCount += loss.sum().item()
|
||||
|
||||
true_preds += (prediction.argmax(dim=1) == labels).sum().item()
|
||||
count += croped_images.size(0)
|
||||
|
||||
val_acc = true_preds / count
|
||||
val_loss = lossCount / count
|
||||
|
||||
if(saving and best_val > val_loss):
|
||||
best_val = val_loss
|
||||
save_dir = os.path.join(SAVE_PATH, model_name)
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
save_path = os.path.join(save_dir, model_name)
|
||||
torch.save(model.state_dict(), save_path)
|
||||
|
||||
print(f"epoch: {epoch+1} | train accuracy: {int(train_acc * 1000) / 10}% | validation accuracy: {int(val_acc * 1000) / 10}% | train loss: {int(train_loss * 1000) / 100} | val loss: {int(val_loss * 1000) / 100}")
|
||||
torch.cuda.empty_cache()
|
||||
return best_val
|
||||
|
||||
|
||||
def trainNormalDataset(model, loss_module, train_loader, val_loader, optimizer, SAVE_PATH, model_name, saving=True, device="cpu"):
|
||||
best_val = torch.finfo(torch.float32).max
|
||||
|
||||
for epoch in range(200):
|
||||
############
|
||||
# Training #
|
||||
############
|
||||
model.train()
|
||||
|
||||
true_preds, count, lossCount = 0, 0, 0.
|
||||
for images, labels in tqdm(train_loader, desc=f"Train", leave=False):
|
||||
images = images.to(device)
|
||||
labels = labels.to(device)
|
||||
|
||||
prediction = model(images)
|
||||
|
||||
loss = loss_module(prediction, labels)
|
||||
lossCount += loss.sum().item()
|
||||
|
||||
optimizer.zero_grad()
|
||||
|
||||
loss.backward()
|
||||
|
||||
optimizer.step()
|
||||
|
||||
true_preds += (prediction.argmax(dim=1) == labels).sum().item()
|
||||
count += images.size(0)
|
||||
|
||||
train_acc = true_preds / count
|
||||
train_loss = lossCount / count
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
##############
|
||||
# Validation #
|
||||
##############
|
||||
model.eval()
|
||||
|
||||
true_preds, count, lossCount = 0, 0, 0.
|
||||
for images, labels in tqdm(val_loader, desc=f"Test", leave=False):
|
||||
with torch.no_grad():
|
||||
images = images.to(device)
|
||||
labels = labels.to(device)
|
||||
|
||||
prediction = model(images)
|
||||
|
||||
loss = loss_module(prediction, labels)
|
||||
lossCount += loss.sum().item()
|
||||
|
||||
true_preds += (prediction.argmax(dim=1) == labels).sum().item()
|
||||
count += images.size(0)
|
||||
|
||||
val_acc = true_preds / count
|
||||
val_loss = lossCount / count
|
||||
|
||||
if(saving and best_val > val_loss):
|
||||
best_val = val_loss
|
||||
save_dir = os.path.join(SAVE_PATH, model_name)
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
save_path = os.path.join(save_dir, model_name)
|
||||
torch.save(model.state_dict(), save_path)
|
||||
|
||||
print(f"epoch: {epoch+1} | train accuracy: {int(train_acc * 1000) / 10}% | validation accuracy: {int(val_acc * 1000) / 10}% | train loss: {int(train_loss * 1000) / 100} | val loss: {int(val_loss * 1000) / 100}")
|
||||
torch.cuda.empty_cache()
|
||||
return best_val
|
||||
|
||||
def resize_keep_aspect(img, target_w, target_h):
|
||||
h, w = img.shape[:2]
|
||||
|
||||
scale = min(target_w / w, target_h / h)
|
||||
new_w = int(w * scale)
|
||||
new_h = int(h * scale)
|
||||
|
||||
resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
|
||||
return resized, w / new_w, h / new_h
|
||||
|
||||
def map_box_to_original(box, scale_x, scale_y):
|
||||
x1, y1, x2, y2 = box
|
||||
|
||||
return int(x1 * scale_x), int(y1 * scale_y), int(x2 * scale_x), int(y2 * scale_y)
|
||||
|
||||
def eval(model, image, BUILD_PATH, device, PADDING = 0, minSize=5, maxSize=600, minConf=0.8):
|
||||
_, H, W = image.shape
|
||||
|
||||
model.load_state_dict(torch.load(BUILD_PATH, map_location=torch.device(device)))
|
||||
model.to(device)
|
||||
model.eval()
|
||||
image_numpy = image.permute(1, 2, 0).cpu().numpy()
|
||||
image_numpy = (image_numpy*255).astype("uint8")
|
||||
image_numpy = cv2.cvtColor(image_numpy, cv2.COLOR_BGR2RGB)
|
||||
image_numpy_copy, resizedW, resizedH = resize_keep_aspect(image_numpy, W, H)
|
||||
|
||||
cv2.imshow("", image_numpy_copy)
|
||||
cv2.waitKey(0)
|
||||
|
||||
ss = cv2.ximgproc.segmentation.createSelectiveSearchSegmentation()
|
||||
ss.setBaseImage(image_numpy_copy)
|
||||
ss.switchToSelectiveSearchFast()
|
||||
|
||||
rects = ss.process()
|
||||
|
||||
draw = image_numpy.copy()
|
||||
predictions = []
|
||||
|
||||
for (x, y, w, h) in tqdm(rects):
|
||||
x, y, w, h = map_box_to_original((x, y, w, h), resizedW, resizedH)
|
||||
|
||||
if w < minSize or h < minSize or w > maxSize or h > maxSize:
|
||||
continue
|
||||
|
||||
x1 = max(x-PADDING, 0)
|
||||
y1 = max(y-PADDING, 0)
|
||||
x2 = min(x+w+PADDING, W-1)
|
||||
y2 = min(y+h+PADDING, H-1)
|
||||
|
||||
crop = image[:, y1:y2, x1:x2]
|
||||
crop = TF.resize(crop, [64, 64], antialias=True)
|
||||
crop = crop.unsqueeze(0).to(device)
|
||||
|
||||
pred = model(crop)
|
||||
probs = F.softmax(pred, dim=1)
|
||||
confidence, cls = torch.max(probs, dim=1)
|
||||
if cls.item() == 1 and confidence.item() > minConf:
|
||||
predictions.append((confidence.item(), cls.item(), (x, y, w, h)))
|
||||
|
||||
predictions = sorted(predictions, key=lambda x: x[0], reverse=True)
|
||||
print(len(predictions))
|
||||
for conf, cls, (x, y, w, h) in predictions[:]:
|
||||
cv2.rectangle(draw, (x, y), (x + w, y + h), (255, 0, 0), 1)
|
||||
|
||||
cv2.imshow("", draw)
|
||||
cv2.waitKey(0)
|
||||
@@ -0,0 +1,56 @@
|
||||
import json
|
||||
import os
|
||||
import cv2
|
||||
from matplotlib import pyplot as plt
|
||||
import numpy as np
|
||||
from torch import tensor
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from torchvision import transforms
|
||||
from torchvision.transforms import ToTensor, functional
|
||||
import torchvision.transforms.functional as TF
|
||||
from tqdm import tqdm
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .r_cnn import ObjectDetectionCNN, train, eval
|
||||
from .cocoDetectionDataset import CocoDetectionDataset
|
||||
import random
|
||||
|
||||
SAVE_PATH = "./saved_models"
|
||||
PADDING = 20
|
||||
|
||||
def get_transform():
|
||||
return ToTensor()
|
||||
|
||||
train_dataset = CocoDetectionDataset(
|
||||
image_dir="data/football/train",
|
||||
annotation_path="data/football/train/_annotations.coco.json",
|
||||
transforms=get_transform()
|
||||
)
|
||||
|
||||
val_dataset = CocoDetectionDataset(
|
||||
image_dir="data/football/valid",
|
||||
annotation_path="data/football/valid/_annotations.coco.json",
|
||||
transforms=get_transform()
|
||||
)
|
||||
|
||||
train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True, collate_fn=lambda x: tuple(zip(*x)))
|
||||
val_loader = DataLoader(val_dataset, batch_size=2, shuffle=True, collate_fn=lambda x: tuple(zip(*x)))
|
||||
|
||||
device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
|
||||
print("Using device", device)
|
||||
|
||||
model = ObjectDetectionCNN(c_in=3, c_hidden=32, c_out=2, layers=10)
|
||||
model.to(device)
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
|
||||
loss_module = nn.CrossEntropyLoss()
|
||||
|
||||
#train(model=model, loss_module=loss_module, train_loader=train_loader, val_loader=val_loader,
|
||||
# optimizer=optimizer, SAVE_PATH=SAVE_PATH, saving=True, PADDING=40, device=device)
|
||||
|
||||
#exit()
|
||||
|
||||
image = next(iter(val_loader))[0][0]
|
||||
eval(model=model, image=image, BUILD_PATH=os.path.join(SAVE_PATH, "object_detection", "object_detection"),
|
||||
device=device, PADDING=40, minSize=5, maxSize=100, minConf=0.8)
|
||||
@@ -0,0 +1,72 @@
|
||||
import os
|
||||
|
||||
from torchvision import transforms
|
||||
from tqdm import tqdm
|
||||
from util import TransformedSubset, visualizeImage
|
||||
from api.src.r_cnn.r_cnn import ObjectDetectionCNN, trainNormalDataset, eval
|
||||
from torch.utils.data import DataLoader, random_split
|
||||
from torchvision import datasets
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torchvision.io import read_image
|
||||
from torch.utils.data import WeightedRandomSampler
|
||||
|
||||
|
||||
|
||||
SAVE_PATH = "./saved_models"
|
||||
PADDING = 20
|
||||
IMAGE_SIZE = 64
|
||||
|
||||
transform = transforms.Compose([
|
||||
transforms.Resize(IMAGE_SIZE),
|
||||
transforms.CenterCrop(IMAGE_SIZE),
|
||||
transforms.ToTensor()
|
||||
])
|
||||
|
||||
transform_augemnt = transforms.Compose([
|
||||
transforms.RandomHorizontalFlip(p=0.5),
|
||||
transforms.RandomRotation(degrees=15),
|
||||
transforms.Resize(IMAGE_SIZE),
|
||||
transforms.CenterCrop(IMAGE_SIZE),
|
||||
transforms.ToTensor()
|
||||
])
|
||||
|
||||
dataset = datasets.ImageFolder("data/faces_processed")
|
||||
|
||||
weights = [1/3, 1]
|
||||
|
||||
sampler = WeightedRandomSampler(
|
||||
weights=weights,
|
||||
num_samples=len(dataset),
|
||||
replacement=True
|
||||
)
|
||||
|
||||
train_size = int(0.8 * len(dataset))
|
||||
val_size = len(dataset) - train_size
|
||||
|
||||
train_subset, val_subset = random_split(dataset, [train_size, val_size])
|
||||
train_dataset = TransformedSubset(train_subset, transform_augemnt)
|
||||
val_dataset = TransformedSubset(val_subset, transform)
|
||||
|
||||
|
||||
train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True)
|
||||
val_loader = DataLoader(val_dataset, batch_size=256, shuffle=False)
|
||||
|
||||
device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
|
||||
print("Using device", device)
|
||||
|
||||
model = ObjectDetectionCNN(c_in=3, c_hidden=32, c_out=2, layers=10)
|
||||
model.to(device)
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
|
||||
loss_module = nn.CrossEntropyLoss(weight=torch.tensor([1/3, 1], device=device, dtype=torch.float32))
|
||||
|
||||
#trainNormalDataset(model=model, loss_module=loss_module, train_loader=train_loader, val_loader=val_loader,
|
||||
# optimizer=optimizer, SAVE_PATH=SAVE_PATH, saving=True, device=device, model_name="face_detection")
|
||||
|
||||
#exit()
|
||||
|
||||
image = read_image("data/faces/train/_url-http_3A_2F_2Fdingyue-ws-126-net_2F2023_2F0201_2F6a17ca80j00rpd7yx00bmd000dw00gop_jpg.rf.Setza3JikTSzG04c0Fd1.jpg")
|
||||
image = image.float() / 255.0
|
||||
visualizeImage(image)
|
||||
|
||||
eval(model=model, image=image, BUILD_PATH=os.path.join(SAVE_PATH, "face_detection", "face_detection"), device=device, minConf=0.97)
|
||||
@@ -0,0 +1,150 @@
|
||||
from enum import Enum
|
||||
import os
|
||||
import sqlite3
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import BaseModel
|
||||
import torch
|
||||
from torchvision import transforms
|
||||
from fastapi import Depends, FastAPI, File, HTTPException, Header, UploadFile
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from PIL import Image
|
||||
import io
|
||||
import torch.nn.functional as F
|
||||
|
||||
from bird_cnn.bird_cnn import Bird_CNN
|
||||
|
||||
import threading
|
||||
|
||||
BUILD_PATH = "../build_models"
|
||||
IMAGE_SIZE = 64
|
||||
|
||||
sem_ai = threading.Semaphore(1)
|
||||
|
||||
class bird_species(Enum):
|
||||
Common_Kingfisher = 0
|
||||
Common_Myna = 1
|
||||
House_Crow = 2
|
||||
Indian_Peacock = 3
|
||||
Indian_Pitta = 4
|
||||
Ruddy_Shelduck = 5
|
||||
Sarus_Crane = 6
|
||||
|
||||
transform = transforms.Compose([
|
||||
transforms.Resize(IMAGE_SIZE),
|
||||
transforms.CenterCrop(IMAGE_SIZE),
|
||||
transforms.ToTensor()
|
||||
])
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
model = Bird_CNN(c_in=3, c_hidden=16, c_out=7)
|
||||
full_path = os.path.join(BUILD_PATH, "bird_cnn")
|
||||
model.load_state_dict(torch.load(full_path, map_location=torch.device(device)))
|
||||
model.to(device)
|
||||
model.eval()
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
origins = [
|
||||
"https://marvinkrausser.com",
|
||||
"https://api.marvinkrausser.com",
|
||||
]
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=origins,
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
@app.post("/predict")
|
||||
async def predict(file: UploadFile = File(...)):
|
||||
with sem_ai:
|
||||
image_bytes = await file.read()
|
||||
|
||||
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
|
||||
image = transform(image).unsqueeze(0).to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
pred = model(image)
|
||||
probs = F.softmax(pred, dim=1)
|
||||
confidence, cls = torch.max(probs, dim=1)
|
||||
|
||||
return {
|
||||
"class": bird_species(cls.item()).name,
|
||||
"confidence": confidence.item()
|
||||
}
|
||||
|
||||
|
||||
load_dotenv("../database/.env")
|
||||
API_KEY = os.getenv("API_KEY")
|
||||
|
||||
def get_api_key(authorization: str = Header(None)):
|
||||
if authorization != f"Bearer {API_KEY}":
|
||||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||||
return authorization
|
||||
|
||||
class Review(BaseModel):
|
||||
website: str
|
||||
rating: int
|
||||
text: str
|
||||
date: str
|
||||
|
||||
sem_db = threading.Semaphore(1)
|
||||
|
||||
@app.get("/review")
|
||||
def get_reviews(auth=Depends(get_api_key)):
|
||||
with sem_db:
|
||||
conn = sqlite3.connect(
|
||||
"database/reviews.db",
|
||||
check_same_thread=False
|
||||
)
|
||||
cur = conn.cursor()
|
||||
|
||||
cur.execute("""
|
||||
SELECT * FROM reviews;
|
||||
""")
|
||||
|
||||
data = cur.fetchall()
|
||||
|
||||
print(data)
|
||||
|
||||
conn.close()
|
||||
return {"data": data}
|
||||
|
||||
@app.post("/review")
|
||||
def post_review(review: Review):
|
||||
with sem_db:
|
||||
conn = sqlite3.connect(
|
||||
"database/reviews.db",
|
||||
check_same_thread=False
|
||||
)
|
||||
cur = conn.cursor()
|
||||
|
||||
cur.execute("""
|
||||
INSERT INTO reviews (website, rating, text, created_at)
|
||||
VALUES (?, ?, ?, ?)
|
||||
""", (review.website, review.rating, review.text, review.date))
|
||||
|
||||
conn.commit()
|
||||
conn.close()
|
||||
return {"status": "ok"}
|
||||
|
||||
@app.delete("/review")
|
||||
def delete_review(auth=Depends(get_api_key)):
|
||||
with sem_db:
|
||||
conn = sqlite3.connect(
|
||||
"database/reviews.db",
|
||||
check_same_thread=False
|
||||
)
|
||||
cur = conn.cursor()
|
||||
|
||||
cur.execute("""
|
||||
DELETE FROM reviews;
|
||||
""")
|
||||
|
||||
conn.commit()
|
||||
conn.close()
|
||||
return {"status": "ok"}
|
||||
@@ -0,0 +1,61 @@
|
||||
from PIL import Image
|
||||
import os
|
||||
from matplotlib import pyplot as plt
|
||||
import torch
|
||||
|
||||
class TransformedSubset(torch.utils.data.Dataset):
|
||||
def __init__(self, subset, transform=None):
|
||||
self.subset = subset
|
||||
self.transform = transform
|
||||
|
||||
def __getitem__(self, idx):
|
||||
x, y = self.subset[idx]
|
||||
|
||||
if self.transform:
|
||||
x = self.transform(x)
|
||||
|
||||
return x, y
|
||||
|
||||
def __len__(self):
|
||||
return len(self.subset)
|
||||
|
||||
|
||||
def visualizeData(dataset):
|
||||
images, labels = next(iter(dataset))
|
||||
|
||||
for i in range(4):
|
||||
img = images[i]
|
||||
|
||||
# Convert tensor shape from [C,H,W] -> [H,W,C]
|
||||
img = img.permute(1, 2, 0)
|
||||
|
||||
plt.figure(figsize=(3,3))
|
||||
plt.imshow(img)
|
||||
plt.title(f"Label: {labels[i].item()}")
|
||||
plt.axis("off")
|
||||
|
||||
plt.show()
|
||||
|
||||
def visualizeImage(image):
|
||||
image = image.permute(1, 2, 0)
|
||||
|
||||
plt.figure(figsize=(3,3))
|
||||
plt.imshow(image)
|
||||
plt.axis("off")
|
||||
|
||||
plt.show()
|
||||
|
||||
def iou(boxA, boxB):
|
||||
xA = max(boxA[0], boxB[0])
|
||||
yA = max(boxA[1], boxB[1])
|
||||
xB = min(boxA[2], boxB[2])
|
||||
yB = min(boxA[3], boxB[3])
|
||||
|
||||
inter_area = max(0, xB - xA) * max(0, yB - yA)
|
||||
|
||||
boxA_area = (boxA[2]-boxA[0]) * (boxA[3]-boxA[1])
|
||||
boxB_area = (boxB[2]-boxB[0]) * (boxB[3]-boxB[1])
|
||||
|
||||
union = boxA_area + boxB_area - inter_area
|
||||
|
||||
return inter_area / union if union > 0 else 0
|
||||
@@ -0,0 +1,55 @@
|
||||
import os
|
||||
|
||||
from torchvision import transforms
|
||||
from tqdm import tqdm
|
||||
from yolo_model import train, Yolo_model
|
||||
from yolo_loss import YoloLoss
|
||||
from cocoDetectionDataset import CocoDetectionDatasetResized
|
||||
from util import TransformedSubset, visualizeImage
|
||||
from torch.utils.data import DataLoader, random_split
|
||||
from torchvision import datasets
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torchvision.io import read_image
|
||||
from torch.utils.data import WeightedRandomSampler
|
||||
|
||||
|
||||
|
||||
SAVE_PATH = "./saved_models"
|
||||
IMAGE_SIZE = 128
|
||||
GRID = 9
|
||||
BATCH_SIZE = 1
|
||||
|
||||
transform = transforms.Compose([
|
||||
transforms.ToTensor()
|
||||
])
|
||||
|
||||
dataset = CocoDetectionDatasetResized(
|
||||
image_dir="data/faces/train",
|
||||
annotation_path="data/faces/train/_annotations.coco.json",
|
||||
img_size=IMAGE_SIZE,
|
||||
transforms=transform
|
||||
)
|
||||
|
||||
train_size = int(0.8 * len(dataset))
|
||||
val_size = len(dataset) - train_size
|
||||
|
||||
train_subset, val_subset = random_split(dataset, [train_size, val_size])
|
||||
|
||||
|
||||
train_loader = DataLoader(train_subset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=lambda x: tuple(zip(*x)))
|
||||
val_loader = DataLoader(val_subset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=lambda x: tuple(zip(*x)))
|
||||
|
||||
device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
|
||||
print("Using device", device)
|
||||
|
||||
model = Yolo_model(c_in=3, c_hidden=32, boxes=2, img_size=IMAGE_SIZE, grid=GRID, labels=1)
|
||||
model.to(device)
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
|
||||
loss_module = YoloLoss()
|
||||
|
||||
train(model=model, loss_module=loss_module, train_loader=train_loader, val_loader=val_loader,
|
||||
optimizer=optimizer, SAVE_PATH=SAVE_PATH, saving=True, device=device,
|
||||
model_name="face_detection_yolo", img_size=IMAGE_SIZE, num_classes=1, grid=GRID)
|
||||
|
||||
exit()
|
||||
@@ -0,0 +1,107 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class YoloLoss(nn.Module):
|
||||
def __init__(self):
|
||||
super(YoloLoss, self).__init__()
|
||||
|
||||
@staticmethod
|
||||
def is_center_in_grid_cell(x, y, img_w, img_h, S, cell_i, cell_j):
|
||||
cell_w = img_w / S
|
||||
cell_h = img_h / S
|
||||
|
||||
# find which cell the center belongs to
|
||||
gt_cell_i = int(y / cell_h)
|
||||
gt_cell_j = int(x / cell_w)
|
||||
|
||||
return (gt_cell_i == cell_i) and (gt_cell_j == cell_j)
|
||||
|
||||
@staticmethod
|
||||
def save_sqrt(i):
|
||||
return torch.sqrt(torch.clamp(i, min=1e-6))
|
||||
|
||||
@staticmethod
|
||||
def iou(boxA, boxB):
|
||||
xA = torch.max(boxA[0], boxB[0])
|
||||
yA = torch.max(boxA[1], boxB[1])
|
||||
xB = torch.min(boxA[2], boxB[2])
|
||||
yB = torch.min(boxA[3], boxB[3])
|
||||
|
||||
inter_area = torch.clamp(xB - xA, min=0) * torch.clamp(yB - yA, min=0)
|
||||
|
||||
boxA_area = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])
|
||||
boxB_area = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])
|
||||
|
||||
union = boxA_area + boxB_area - inter_area
|
||||
|
||||
return inter_area / (union + 1e-6)
|
||||
|
||||
@staticmethod
|
||||
def xywh_to_xyxy(box):
|
||||
x, y, w, h = box
|
||||
|
||||
return torch.stack([
|
||||
x - w / 2,
|
||||
y - h / 2,
|
||||
x + w / 2,
|
||||
y + h / 2
|
||||
])
|
||||
|
||||
def forward(self, pred, target, labels, cell_size, img_w, img_h, factor_no_object = 0.5, factor_object = 5):
|
||||
#pred: [SxSx(B(x,y,w,h,con)C)]
|
||||
#target: [B(x, y, w, h, C)]
|
||||
|
||||
for s0 in range(pred.shape[0]):
|
||||
for s1 in range(pred.shape[1]):
|
||||
boxes = pred[s0, s1, :(-1*labels)]
|
||||
pred_label = pred[s0, s1, -1:]
|
||||
|
||||
ground_truth_box = None
|
||||
for b in range(target.shape[0] // (4 + labels)):
|
||||
box = target[(b)*(4 + labels):(b+1)*(4 + labels)]
|
||||
if self.is_center_in_grid_cell(box[0]+box[2]/2, box[1]+box[3]/2, img_w, img_h, cell_size, s0, s1):
|
||||
ground_truth_box = box[:4]
|
||||
ground_truth_label = box[4:]
|
||||
break
|
||||
|
||||
|
||||
coord_loss = torch.tensor(0., device=pred.device)
|
||||
size_loss = torch.tensor(0., device=pred.device)
|
||||
prob_loss = torch.tensor(0., device=pred.device)
|
||||
label_loss = torch.tensor(0., device=pred.device)
|
||||
|
||||
if ground_truth_box is not None:
|
||||
|
||||
iou_stats = []
|
||||
|
||||
for b in range(boxes.shape[0] // 5):
|
||||
box = boxes[(b)*5:(b+1)*5]
|
||||
|
||||
iou_stats.append(self.iou(self.xywh_to_xyxy(box[0:4]), self.xywh_to_xyxy(ground_truth_box[0:4])))
|
||||
iou_tensor = torch.stack(iou_stats)
|
||||
responsible = torch.argmax(iou_tensor)
|
||||
|
||||
for b in range(boxes.shape[0] // 5):
|
||||
box = boxes[(b)*5:(b+1)*5]
|
||||
|
||||
is_responsible = (b == responsible).float()
|
||||
|
||||
coord_loss += factor_object * is_responsible * ((box[0] - ground_truth_box[0])**2 + (box[1] - ground_truth_box[1])**2)
|
||||
size_loss += factor_object * is_responsible * ((self.save_sqrt(box[2]) - self.save_sqrt(ground_truth_box[2]))**2 +
|
||||
(self.save_sqrt(box[3]) - self.save_sqrt(ground_truth_box[3]))**2)
|
||||
|
||||
prob_loss += is_responsible * (box[4] - 1)**2
|
||||
prob_loss += factor_no_object * (1 - is_responsible) * box[4]**2
|
||||
|
||||
|
||||
for i in range(pred_label.shape[0]):
|
||||
label_loss += ((pred_label[i] - ground_truth_label[i])**2) / pred_label.shape[0]
|
||||
|
||||
|
||||
else:
|
||||
for b in range(boxes.shape[0] // 5):
|
||||
box = boxes[(b)*5:(b+1)*5]
|
||||
|
||||
prob_loss += factor_no_object * box[4] ** 2
|
||||
|
||||
return (label_loss + prob_loss + coord_loss + size_loss) / pred.shape[0]**2
|
||||
@@ -0,0 +1,134 @@
|
||||
import os
|
||||
import cv2
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torchvision.transforms.functional as TF
|
||||
import time
|
||||
|
||||
|
||||
class Yolo_model(nn.Module):
|
||||
def __init__(self, c_in, c_hidden, boxes, img_size, grid, labels):
|
||||
super().__init__()
|
||||
|
||||
self.grid = grid
|
||||
|
||||
self.model = nn.Sequential(
|
||||
nn.Conv2d(in_channels=c_in, out_channels=c_hidden, kernel_size=7, padding=3),
|
||||
nn.LeakyReLU(),
|
||||
nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(),
|
||||
nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(),
|
||||
nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(),
|
||||
nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(),
|
||||
nn.Conv2d(in_channels=c_hidden, out_channels=c_hidden, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(),
|
||||
nn.Flatten(),
|
||||
nn.Linear(in_features=img_size*img_size*c_hidden, out_features=grid*grid*(boxes*5+labels))
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
batch = x.shape[0]
|
||||
x = self.model(x)
|
||||
x = x.reshape(batch, self.grid, self.grid, -1)
|
||||
return x
|
||||
|
||||
def create_stack(images, annotations, num_classes, device):
|
||||
data = []
|
||||
labels = []
|
||||
for i in range(len(images)):
|
||||
image = images[i]
|
||||
|
||||
annotation = annotations[i]
|
||||
|
||||
label_train = []
|
||||
for box, label in tuple(zip(annotation["boxes"], annotation["labels"])):
|
||||
box = box
|
||||
label = 0
|
||||
one_hot = F.one_hot(torch.tensor(label, dtype=torch.long), num_classes=num_classes)
|
||||
label_train.append(torch.cat((box, one_hot), dim=0))
|
||||
|
||||
label_torch = torch.stack(label_train).reshape(-1).to(device)
|
||||
labels.append(label_torch)
|
||||
data.append(image)
|
||||
|
||||
return data, labels
|
||||
|
||||
|
||||
def train(model, num_classes, loss_module, train_loader, val_loader, optimizer, SAVE_PATH, model_name, grid, saving=True, device="cpu", img_size=64):
|
||||
best_val = torch.finfo(torch.float32).max
|
||||
|
||||
for epoch in range(200):
|
||||
############
|
||||
# Training #
|
||||
############
|
||||
model.train()
|
||||
|
||||
count, lossCount = 0, 0.
|
||||
for images, annotations in tqdm(train_loader, desc=f"Train", leave=False):
|
||||
data, labels = create_stack(images, annotations, num_classes, device)
|
||||
|
||||
data = torch.stack(data).to(device)
|
||||
|
||||
prediction = model(data)
|
||||
|
||||
loss = []
|
||||
for pred, label in tuple(zip(prediction, labels)):
|
||||
loss.append(loss_module(pred = pred, target= label, labels=num_classes, cell_size=grid, img_w = img_size, img_h = img_size))
|
||||
loss = torch.stack(loss).to(device)
|
||||
loss = loss.sum()
|
||||
lossCount += loss.item()
|
||||
|
||||
optimizer.zero_grad()
|
||||
|
||||
loss.backward()
|
||||
|
||||
optimizer.step()
|
||||
|
||||
count += data.size(0)
|
||||
|
||||
train_loss = lossCount / count
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
##############
|
||||
# Validation #
|
||||
##############
|
||||
model.eval()
|
||||
|
||||
count, lossCount = 0, 0.
|
||||
for images, annotations in tqdm(val_loader, desc=f"Test", leave=False):
|
||||
with torch.no_grad():
|
||||
data, labels = create_stack(images, annotations, num_classes, device)
|
||||
|
||||
data = torch.stack(data).to(device)
|
||||
|
||||
prediction = model(data)
|
||||
|
||||
loss = []
|
||||
for pred, label in tuple(zip(prediction, labels)):
|
||||
loss.append(loss_module(pred = pred, target= label, labels=num_classes, cell_size=grid, img_w = img_size, img_h = img_size))
|
||||
loss = torch.stack(loss).to(device)
|
||||
|
||||
|
||||
lossCount += loss.sum().item()
|
||||
|
||||
count += data.size(0)
|
||||
|
||||
val_loss = lossCount / count
|
||||
|
||||
if(saving and best_val > val_loss):
|
||||
best_val = val_loss
|
||||
save_dir = os.path.join(SAVE_PATH, model_name)
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
save_path = os.path.join(save_dir, model_name)
|
||||
torch.save(model.state_dict(), save_path)
|
||||
|
||||
print(f"epoch: {epoch+1} | train loss: {int(train_loss * 1000) / 100} | val loss: {int(val_loss * 1000) / 100}")
|
||||
torch.cuda.empty_cache()
|
||||
return best_val
|
||||
Reference in New Issue
Block a user