Azeem Aslam 👋

import torch, cv2, numpy as np, torch.nn as nn
from ultralytics import YOLO
from torchvision.models import resnet50

model = YOLO("yolov9.pt")  # load pretrained detection weights
results = model.predict(source=img, conf=0.25, iou=0.45)

for r in results:
    boxes = r.boxes.xyxy       # bounding box coordinates
    scores = r.boxes.conf      # confidence scores per detection
    classes = r.boxes.cls      # predicted class indices

def preprocess(image, size=(640, 640)):
    image = cv2.resize(image, size)
    tensor = torch.from_numpy(image).float()
    return tensor.permute(2, 0, 1) / 255.0

class VisionModel(nn.Module):
    def __init__(self, num_classes=80):
        super().__init__()
        self.backbone = resnet50(pretrained=True)
        self.head = nn.Linear(2048, num_classes)

    def forward(self, x):
        features = self.backbone(x)
        return self.head(features)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-5)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(epochs):
    loss = train_step(model, loader, optimizer, loss_fn)
    print(f"epoch={epoch:03d} loss={loss:.4f} lr={optimizer.param_groups[0]['lr']:.6f}")

def evaluate(model, val_loader):
    model.eval()
    correct, total = 0, 0
    for images, labels in val_loader:
        outputs = model(images)
        preds = outputs.argmax(dim=1)
        correct += (preds == labels).sum().item()
        total += labels.size(0)
    return correct / total  # returns overall validation accuracy

model.save("checkpoints/best.pt")
print("training complete. accuracy:", evaluate(model, val_loader))