changed api structure and added object detection
This commit is contained in:
@@ -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)
|
||||
Reference in New Issue
Block a user