yingkunwu / r-yolov4 Goto Github PK

# 官方示例代码
from torch.utils.tensorboard import SummaryWriter
import numpy as np

writer = SummaryWriter()

for n_iter in range(100):
    writer.add_scalar('Loss/train', np.random.random(), n_iter)

<class_id> <centerX/imgsz> <centerY/imgsz> <w/imgsz> <h/imgsz> <theta/360 or 2P?>

import os, sys
currentdir = os.path.dirname(os.path.realpath(__file__))
parentdir = os.path.dirname(currentdir)
sys.path.append(parentdir)

from label_studio_ml.model import LabelStudioMLBase
from label_studio_ml.utils import get_image_size, get_single_tag_keys, is_skipped
from label_studio.core.utils.io import json_load, get_data_dir 
from label_studio.core.settings.base import DATA_UNDEFINED_NAME

import time
import random
import numpy as np
import torch
import shutil
import json
from terminaltables import AsciiTable
import glob

from model.yolo import Yolo
from lib.utils import load_class_names
from lib.scheduler import CosineAnnealingWarmupRestarts
from lib.post_process import post_process
from lib.logger import *
from lib.options import LabelStudioOptions
from lib.plot import rescale_boxes
import label_studio_sdk
from datasets.label_studio_dataset import ImageDataset, LabelStudioDataset, get_transformed_image
import cv2 as cv

from urllib.parse import urlparse

from PIL import Image

print("LabelStudioSdk Version: ", label_studio_sdk.__version__)

LABEL_STUDIO_HOST = os.getenv('LABEL_STUDIO_HOST', 'http://localhost:8080')
LABEL_STUDIO_API_KEY = os.getenv('LABEL_STUDIO_API_KEY', '4c23feec13e2118e053b9a9940f73ed96c0e0841')

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

def weights_init_normal(m):
    if isinstance(m, torch.nn.Conv2d):
        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif isinstance(m, torch.nn.BatchNorm2d):
        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
        torch.nn.init.constant_(m.bias.data, 0.0)

def init():
    random.seed(42)
    np.random.seed(42)
    torch.manual_seed(42)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

class RotBBoxModel(object):
    def __init__(self, num_classes, args):
        self.args = args
        
        self.model = Yolo(n_classes=num_classes)
        self.model = self.model.to(device)

        self.logger = None
        self.model_path = None

        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.args.lr)

    def log(self, total_loss, num_epochs, epoch, global_step, total_step, start_time):
        log = "\n---- [Epoch %d/%d] ----\n" % (epoch + 1, num_epochs)

        tensorboard_log = {}
        loss_table_name = ["Step: %d/%d" % (global_step, total_step),
                            "loss", "reg_loss", "conf_loss", "cls_loss"]
        loss_table = [loss_table_name]

        temp = ["YoloLayer1"]
        for name, metric in self.model.yolo1.metrics.items():
            if name in loss_table_name:
                temp.append(metric)
            tensorboard_log[f"{name}_1"] = metric
        loss_table.append(temp)

        temp = ["YoloLayer2"]
        for name, metric in self.model.yolo2.metrics.items():
            if name in loss_table_name:
                temp.append(metric)
            tensorboard_log[f"{name}_2"] = metric
        loss_table.append(temp)

        temp = ["YoloLayer3"]
        for name, metric in self.model.yolo3.metrics.items():
            if name in loss_table_name:
                temp.append(metric)
            tensorboard_log[f"{name}_3"] = metric
        loss_table.append(temp)

        tensorboard_log["total_loss"] = total_loss
        self.logger.list_of_scalars_summary(tensorboard_log, global_step)

        log += AsciiTable(loss_table).table
        log += "\nTotal Loss: %f, Runtime: %f\n" % (total_loss, time.time() - start_time)
        print(log)

    def save(self, path):
        print("Model saved in: ", path)
        torch.save(self.model.state_dict(), path)

    def load(self, path, train=False):
        print("Loading model...")
        if not train:
            print("Loading model for prediction...")
            self.model_path = path
            if os.path.exists(self.model_path):
                weight_path = glob.glob(os.path.join(self.model_path, "*.pth"))
                if len(weight_path) == 0:
                    assert False, "Model weight not found"
                elif len(weight_path) > 1:
                    assert False, "Multiple weights are found. Please keep only one weight in your model directory"
                else:
                    weight_path = weight_path[0]
            else:
                assert False, "Model is not exist"
            pretrained_dict = torch.load(weight_path, map_location=device)
            self.model.load_state_dict(pretrained_dict)
            self.model.eval()
        else:
            print("Loading model for training...")
            # if os.path.exists(path):
            #     weight_path = glob.glob(os.path.join(path, "*.pth"))[0]
            #     print("weight_path: ", weight_path)
            # else:
            #     print("Path does not exist")
            weight_path = "weights/pretrained/yolov4.pth"
            pretrained_dict = torch.load(weight_path, map_location=device)
            model_dict = self.model.state_dict()

            # 1. filter out unnecessary keys
            # pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
            pretrained_dict = {k: v for i, (k, v) in enumerate(pretrained_dict.items()) if i < 552}
            # 2. overwrite entries in the existing state dict
            model_dict.update(pretrained_dict)
            # 3. load the new state dict
            self.model.apply(weights_init_normal)
            self.model.load_state_dict(model_dict)
            self.model.eval()

    def predict(self, image_urls):
        images = torch.stack([get_transformed_image(url, self.args.img_size) for url in image_urls]).to(device)
        
        with torch.no_grad():
            temp = time.time()
            output, _ = self.model(images)  # batch=1 -> [1, n, n], batch=3 -> [3, n, n]
            temp1 = time.time()
            boxes = post_process(output, self.args.conf_thres, self.args.nms_thres)
            temp2 = time.time()
            print('-----------------------------------')
            num = 0
            for b in boxes:
                if b is None:
                    break
                num += len(b)
            print("{} objects found".format(num))
            print("Inference time : ", round(temp1 - temp, 5))
            print("Post-processing time : ", round(temp2 - temp1, 5))
            print('-----------------------------------')
            return boxes

    def train(self, dataloader, num_epochs=5):
        init()
        if(self.model_path == None):
            self.model_path = os.path.join("weights", self.args.model_name)
        self.logger = Logger(os.path.join(self.model_path, "logs"))

        num_iters_per_epoch = len(dataloader)
        scheduler_iters = round(num_epochs * len(dataloader) / self.args.subdivisions)
        total_step = num_iters_per_epoch * num_epochs

        scheduler = CosineAnnealingWarmupRestarts(self.optimizer,
                                                first_cycle_steps=round(scheduler_iters),
                                                max_lr=self.args.lr,
                                                min_lr=1e-5,
                                                warmup_steps=round(scheduler_iters * 0.1),
                                                cycle_mult=1,
                                                gamma=1)

        start_time = time.time()
        self.model.train()
        for epoch in range(num_epochs):
            print('Epoch {}/{}'.format(epoch, num_epochs - 1))
            print('-' * 10)

            for batch, (_, imgs, targets) in enumerate(dataloader):
                global_step = num_iters_per_epoch * epoch + batch + 1
                imgs = imgs.to(device)
                targets = targets.to(device)

                outputs, loss = self.model(imgs, targets)

                loss.backward()
                total_loss = loss.detach().item()

                if global_step % self.args.subdivisions == 0:
                    self.optimizer.step()
                    self.optimizer.zero_grad()
                    scheduler.step()

                self.log(total_loss, num_epochs, epoch, global_step, total_step, start_time)

        print()

        time_elapsed = time.time() - start_time
        print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))

        return self.model

class RotBBoxModelApi(LabelStudioMLBase):
    
    def __init__(self, **kwargs):
        # don't forget to initialize base class...
        super(RotBBoxModelApi, self).__init__(**kwargs)
        
        parser = LabelStudioOptions()
        self.args = parser.parse()
        
        self.from_name, self.to_name, self.value, self.labels_in_config = get_single_tag_keys(
            self.parsed_label_config, 'RectangleLabels', 'Image'
        )

        print("from_name: ", self.from_name)
        print("to_name: ", self.to_name)
        print("value: ", self.value)
        print("labels_in_config: ", self.labels_in_config)
        print("parsed_label_config: ", self.parsed_label_config)
        print("train_output: ", self.train_output)

        # self.model = RotBBoxModel(len(self.labels_in_config), self.args)
        # self.model_path = os.path.join("weights", self.args.model_name)
        # print(self.model_path)
        # self.model.load(self.model_path)

        if self.train_output:
            self.model = RotBBoxModel(len(self.labels_in_config), self.args)
            self.model.load(self.train_output['model_path'], self.train_output)
        else:
            self.model = RotBBoxModel(len(self.labels_in_config), self.args)
            model_path = os.path.join("weights", self.args.model_name)
            print(model_path)
            self.model.load(model_path)

    def reset_model(self):
        self.model = RotBBoxModel(len(self.labels_in_config), self.args)
        self.model_path = os.path.join("weights", self.args.model_name)
        self.model.load(self.model_path)
   
    def predict(self, tasks, **kwargs):
        """ This is where inference happens:
            model returns the list of predictions based on input list of tasks

            :param tasks: Label Studio tasks in JSON format
        """
        image_urls = [task['data'][self.value] for task in tasks]
        print(image_urls)
        model_results = self.model.predict(image_urls)
        results = []
        all_scores = []
        avg_score = 0

        for i, (url, box) in enumerate(zip(image_urls, model_results)):
            if box is not None:
                image_path = self.get_local_path(url)
                # image_shape = get_image_shape(url)
                img_width, img_height = get_image_size(image_path)
                boxes = rescale_boxes(box, self.args.img_size, (img_height, img_width))
                boxes = np.array(boxes)

                for i in range(len(boxes)):
                    bbox = boxes[i]
                    center_x, center_y, w, h, theta = bbox[0], bbox[1], bbox[2], bbox[3], bbox[4]
                    score = round(bbox[5] * bbox[6], 2)
                    cls_id = np.squeeze(int(bbox[7]))

                    # Calculate top left corner of rotated bbox (box center is origin)
                    left_local = -w/2
                    top_local = -h/2
                    rotated_left_local = np.cos(theta) * left_local - np.sin(theta) * top_local
                    rotated_top_local = np.sin(theta) * left_local + np.cos(theta) * top_local
                    rotated_left = center_x + rotated_left_local
                    rotated_top = center_y + rotated_top_local

                    x_percent = ( (rotated_left / img_width) * 100.0).item()
                    y_percent = ( (rotated_top / img_height) * 100.0).item()
                    w_percent = ( (w / img_width) * 100.0).item()
                    h_percent = ( (h / img_height) * 100.0).item()

                    results.append({
                        'from_name': self.from_name,
                        'to_name': self.to_name,
                        'type': 'rectanglelabels',
                        'value': {
                            'rectanglelabels': [self.labels_in_config[cls_id]],
                            'x': x_percent,
                            'y': y_percent,
                            'width': w_percent,
                            'height': h_percent,
                            'rotation': np.rad2deg(theta).item()
                        },
                        'score': score.item()
                    })
                    all_scores.append(score)
                avg_score = sum(all_scores) / max(len(all_scores), 1)
        if(avg_score != 0):
            avg_score = avg_score.item()
        return [{
            'result': results,
            'score': avg_score
        }]

    def download_tasks(self, project):
        """
        Download all labeled tasks from project using the Label Studio SDK.
        Read more about SDK here https://labelstud.io/sdk/
        :param project: project ID
        :return:
        """
        ls = label_studio_sdk.Client(LABEL_STUDIO_HOST, LABEL_STUDIO_API_KEY)
        project = ls.get_project(id=project)
        tasks = project.get_labeled_tasks()
        return tasks

    def fit(self, tasks, workdir=None, batch_size=32, num_epochs=10, **kwargs):
        """
        This method is called each time an annotation is created or updated
        :param kwargs: contains "data" and "event" key, that could be used to retrieve project ID and annotation event type
                        (read more in https://labelstud.io/guide/webhook_reference.html#Annotation-Created)
        :return: dictionary with trained model artefacts that could be used further in code with self.train_output
        """
        if 'data' not in kwargs:
            raise KeyError(f'Project is not identified. Go to Project Settings -> Webhooks, and ensure you have "Send Payload" enabled')
        
        data = kwargs['data']
        project = data['project']['id']
        tasks = self.download_tasks(project)
        if len(tasks) > 0:
            print(f'{len(tasks)} labeled tasks downloaded for project {project}')
            
            image_urls, image_labels = [], []
            print('Collecting annotations...')
            for task in tasks:
                
                if is_skipped(task):
                    continue

                filepath = self.get_local_path(task['data'][self.value])
                image_urls.append(filepath)
                image_labels.append(task['annotations'][0]['result'])


            # augment = False if self.args.no_augmentation else True
            # mosaic = False if self.args.no_mosaic else True
            # multiscale = False if self.args.no_multiscale else True

            augment = False
            mosaic = False
            multiscale = False

            print(f'Creating dataset with {len(image_urls)} images...')
            dataset = LabelStudioDataset(image_urls, image_labels, self.labels_in_config, 
                                         self.args.img_size, self.args.sample_size,
                                         augment=augment, mosaic=mosaic, multiscale=multiscale)
            dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True, pin_memory=True, collate_fn=dataset.collate_fn)
            
            print('Train model...')
            self.reset_model()
            self.model.train(dataloader, num_epochs=num_epochs)

            print('Save model...')
            # model_path = os.path.join(workdir, 'model.pt')
            model_path = os.path.join(self.model_path, "ryolov4.pth")
            self.model.save(model_path)

            return {
                'model_path': model_path, 
                'labels': image_labels
            }

        else:
            print('No labeled tasks found: make some annotations...')
            return {}

 44 def skewiou(box1, box2):
 45 

How can I transfer learning of some specific classes from the pre-trained weights (e.g. dota) to my custom training if my custom classes are different from the pre-trained classes?

python train.py --data ../dataset/dataset.yaml --config /home/prefect/Work/PalmDetector/Model/OBB/R-YOLOv4/data/hyp.yaml --img_size 640 --epochs 300

Namespace(batch_size=4, config='/home/prefect/Work/PalmDetector/Model/OBB/R-YOLOv4/data/hyp.yaml', data='../dataset/dataset.yaml', epochs=300, img_size=640, lr=0.01, mode='csl', model_name='trash', optimizer='SGD', ver='yolov5', weights_path='')
2023-12-13 00:56:13 WARNING  Model name exists, do you want to override the previous model?
>> [Y:N]y
Traceback (most recent call last):
  File "train.py", line 266, in <module>
    t.train()
  File "train.py", line 145, in train
    train_dataset, train_dataloader = load_data(
  File "/home/prefect/Work/PalmDetector/Model/OBB/R-YOLOv4/lib/load.py", line 15, in load_data
    dataset = CustomDataset(data_dir, class_names, hyp, img_size=640, augment=augment, csl=csl)
TypeError: __init__() got multiple values for argument 'img_size'